junhai0428/yolov5_obb_kld

KLD实现旋转目标检测

YOLOv5_OBB_KLD

基于kld_loss的YOLOv5 旋转目标检测。因为基于项目 也可以YOLOv5 in DOTA_OBB dataset with CSL_label.(Oriented Object Detection)

Datasets and pretrained checkpoint

• Datasets : DOTA
• Pretrained Checkpoint or Demo Files :
  • train,detect_and_evaluate_demo_files: | Baidu Drive(6666). | Google Drive |
  • yolov5x.pt: | Baidu Drive(6666). | Google Drive |
  • yolov5l.pt: | Baidu Drive(6666). | Google Drive |
  • yolov5m.pt: | Baidu Drive(6666). | Google Drive |
  • yolov5s.pt: | Baidu Drive(6666). | Google Drive |
  • YOLOv5_DOTAv1.5_OBB.pt: | Baidu Drive(6666). | Google Drive |

Fuction

• train.py. Train.

• detect.py. Detect and visualize the detection result. Get the detection result txt.

• evaluation.py. Merge the detection result and visualize it. Finally evaluate the detector

Installation (Linux Recommend, Windows not Recommend)

1. Python 3.8 with all requirements.txt dependencies installed, including torch==1.6, opencv-python==4.1.2.30, To install run:

$   pip install -r requirements.txt

2. Install swig

$   cd  \.....\yolov5_OBB_KLD\utils
$   sudo apt-get install swig

3. Create the c++ extension for python

$   swig -c++ -python polyiou.i
$   python setup.py build_ext --inplace

效果

只是用DOTAv1.5的ship一类进行训练,超参数相同,KLD比CSL的AP50高0.3%,不过收敛很快。

More detailed explanation

想要了解其他相关实现的细节和原理可以参看项目 的README.md 这里主要介绍修改成KLD_LOSS的部分。 1. 'train.py'

• parser.add_argument('--use_kld', type=bool, default=True, help='use kld')选择KLD or CSL
• 修改.\models\yolo.py的Detect类中初始化函数的self.angle = 1 #CSL---180 KLD--1

2. 'test.py'

• 新增了在线推断代码

3. 'detect.py'

• 新增了多batch_size的检测,修改ManyPi=True
• parser.add_argument('--kld', type=bool, default=True, help='use kld') 对应KLD or CSL的检测

4. 'evaluation.py'

• Run the detect.py demo first. Then change the path with yours:
• 添加了merged前后的预测结果.

5. 'yolo_new.py'

• 在YOLOv5中加入注意力机制

感激

感谢以下的项目,排名不分先后

• BossZard/rotation-yolov5
• hukaixuan19970627/YOLOv5_DOTA_OBB.
• SJTU-Thinklab-Det/DOTA-DOAI
• buzhidaoshenme/YOLOX-OBB

关于作者

  Name  : "lx"
  describe myself："good man"