【目标检测】ATSS复现

这是一个目标检测中选择正负样本的策略。

1.环境准备

ubuntu16.04
cuda10.1
cudnn7
python3.6

           

git clone https://github.com/sfzhang15/ATSS.git
cd ATSS

# because of cuda version error: no AT_CHECK defination
sed -i '16 a #ifndef AT_CHECK \n#define AT_CHECK TORCH_CHECK \n#endif' atss_core/csrc/cuda/deform_pool_cuda.cu
sed -i '13 a #ifndef AT_CHECK \n#define AT_CHECK TORCH_CHECK \n#endif' atss_core/csrc/cuda/deform_conv_cuda.cu
           

2.模型准备

wget https://drive.google.com/u/0/uc?export=download&confirm=xa89&id=1t8RLdQ6fsFXa0kzPIQ7541uZeQeMXP73
           

3.demo运行

修改：

vim ./ATSS/atss_core/modeling/rpn/atss/inference.py +46

# 修改：
pre_nms_top_n = candidate_inds.view(N, -1).sum(1)

#为：
pre_nms_top_n = candidate_inds.reshape(N, -1).sum(1)
           

想要保存结果的话，可以在demo/atss_demo.py的104行添加如下代码：

save_path = os.path.join(args.images_dir, "save")
if not os.path.exists(save_path):
    os.makedirs(save_path)
cv2.imwrite(os.path.join(save_path, im_name), composite)
           

运行demo： 

python demo/atss_demo.py
           

结果如下：

 4.后续

后续有时间的话，看下怎么加到我们自己的模型中去训练。