【目标檢測】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.後續

後續有時間的話，看下怎麼加到我們自己的模型中去訓練。