yolov5_obb的兼容性改进

最近在学习yolo的旋转框检测，选用了yolov5_obb进行测试，结果代码由于各种问题跑不起来。我将里面用到的废弃函数进行了更改，以便能用用新环境进行运行。

我将修改完的代码上传到了我的Github仓库里。

改进方法

utils/nms_rotated/nms_rotated_wrapper.py

运行的时候会显示too_small和inds这两个变量和ori_inds不在同一个设备上，需要添加代码将两个变量移动到ori_inds所在的设备上。

ori_inds = torch.arange(dets_th.size(0)) # 0 ~ n-1
# add this
too_small = too_small.to(ori_inds.device)
# add end
ori_inds = ori_inds[~too_small]
dets_th = dets_th[~too_small] # (n_filter, 5)
scores = scores[~too_small]

inds = nms_rotated_ext.nms_rotated(dets_th, scores, iou_thr)
# add this
inds = inds.to(ori_inds.device)
# add end
inds = ori_inds[inds]

utils/nms_rotated/src/poly_nms_cuda.cu

THC库已经被弃用，我们需要用ATen库进行更换

这里我是根据这个博客进行更改的

删除#include <THC/THC.h>头文件，把所有的THCudaCheck(cudaGetLastError());替换成 AT_CUDA_CHECK(cudaGetLastError());

对每个.cu文件中调用THCCeilDiv(x,y)的地方，都把这个函数换成 (x+y-1)/y的形式，如：

1 2	// const int col_blocks = THCCeilDiv(n_polys, threadsPerBlock); const int col_blocks = (n_polys+threadsPerBlock-1) / threadsPerBlock;

添加头文件#include <ATen/cuda/ThrustAllocator.h>

注释掉THCState *state = at::globalContext().lazyInitCUDA();

修改THCudaMalloc函数：

1
2

// mask_dev = (unsigned long long*) THCudaMalloc(state, boxes_num * col_blocks * sizeof(unsigned long long));
mask_dev = (unsigned long long*) c10::cuda::CUDACachingAllocator::raw_alloc(boxes_num * col_blocks * sizeof(unsigned long long));

替换THCudaFree

1 2	// THCudaFree(state, mask_dev); c10::cuda::CUDACachingAllocator::raw_delete(mask_dev);

将THCudaCheck替换成C10_CUDA_CHECK

将AT_CHECK替换为TORCH_CHECK

utils/datasets.py

将所有的np.int改写为int

utils/general.py

torch.load函数添加weights_only=False参数

utils/loss.py

这个文件会报错说数据类型转换错误。需要先将变量转换为long，注意这里不是原地转换，需要用赋值的形式进行转换。

# indices.append((b, a, gj.clamp_(0, feature_wh[1] - 1), gi.clamp_(0, feature_wh[0] - 1)))  # image, anchor, grid indices
gi = gi.clamp(0, feature_wh[0] - 1).long()
gj = gj.clamp(0, feature_wh[1] - 1).long()
indices.append((b, a, gj, gi))  # image, anchor, grid indices

utils/plots.py

self.font.getsize函数已经被废弃。

1
2
3

# w, h = self.font.getsize(label)  # text width, height
bbox = self.font.getbbox(label)
w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]

models/experimental.py

torch.load函数添加weights_only=False参数

train.py

torch.load函数添加weights_only=False参数, 注释掉下面两行防止更新numpy。

1 2	# check_git_status() # check_requirements(exclude=['thop'])

编译C++代码

原作者用C++写了一些函数，需要编译后才能使用。但是我按照他写的命令缺一直报错说ModuleNotFoundError: No module named 'torch'，显然我的环境里是有torch的，但是编译的时候报错说没有。

非常玄学的是，我确实有一个环境能够通过python setup.py develop命令来进行编译的，我记得当时我安装了一个pycocotools库之后就可以了。但是过了一天，我想从头到尾再创建一个新的虚拟环境过一遍的时候发现又不能编译了，我查了一下pycocotools仅仅是python对coco数据集做的小工具，应该不是我编译成功的关键要素。我也尝试过sudo apt install python-dev安装编译工具，但是没有用。我导出pip requirement和conda list并且仔细对比了一下发现里面的东西都是一样的。我寻找了很多网页都没有发现是怎么回事，有一些人和我一样遇到了一样的问题，但是都没有解决方案。

我觉得可能是代码与新环境不匹配的问题。我在和chatgpt讨论如何更改代码的时候，他提到了一个新的执行命令：