YOLOv5_OBB_KLD
YOLOv5_OBB_KLD copied to clipboard
KLD实现旋转目标检测
YOLOv5_OBB_KLD



代码实现了基于YOLOv5的遥感旋转框检测。利用CSL和KLD实现角度的学习,并且加入注意力机制提高检测效果。
数据集和权重文件
-
原始预训练的权重文件
:-
yolov5x.pt、yolov5l.pt、yolov5m.pt、yolov5s.pt
: Baidu Drive(6666).
-
-
训练好的部分权重文件
:-
YOLOv5_DOTAv1.5_OBB.pt
: Baidu Drive(6666)
-
项目安装 (支持Linux系统)
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
训练
-
train.py
. Note:修改.\models\yolo.py
的Detect类
中初始化函数的self.angle = 180 #CSL对应180 KLD对应1
,默认使用CSL.
python train.py --weights weights/yolov5m.pt --cfg models/yolov5m.yaml --use_kld False --device 0 --epochs 300 --batch_size 4 --workers 4 --logdir runs/
评估
-
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
python detect.py --weights runs/exp/weights/best.pt --source 'dataset path' --output 'output path' --conf_thres 0.35 --iou_thres 0.4 --device 0 --kld False
python evaluation.py
''' example
检测结果已merge
检测结果已按照类别分类
校验数据集名称文件已生成
classname: ship
P: 0.8550878121966288
R: 0.900046446818393
[email protected]: 0.8889719225631516
classaps: [ 88.897]
原始存在文件,删除
检测结果已按照类别分类
校验数据集名称文件已生成
classname: ship
P: 0.8511538986754063
R: 0.8677432827509397
[email protected]: 0.8096364184338725
classaps: [ 80.964]
'''
结果展示
数据集图片尺寸裁剪为1024*1024,gap为10%。实验中NMS时统一使用的置信度阈值是0.35,IoU阈值是0.4。

感激
感谢以下的项目,排名不分先后
- BossZard/rotation-yolov5
- hukaixuan19970627/YOLOv5_DOTA_OBB.
- SJTU-Thinklab-Det/DOTA-DOAI
- buzhidaoshenme/YOLOX-OBB
关于作者
Name : "lx"
describe myself:"good man"