├── README.md ├── yolov5_obb ├── Arial.ttf ├── DOTA_devkit │ ├── DOTA.py │ ├── DOTA2COCO.py │ ├── DOTA2JSON.py │ ├── ImgSplit.py │ ├── ImgSplit_multi_process.py │ ├── ResultEnsembleNMS_multi_process.py │ ├── ResultMerge.py │ ├── ResultMerge_multi_process.py │ ├── SplitOnlyImage.py │ ├── SplitOnlyImage_multi_process.py │ ├── Visuallize_dota.py │ ├── __init__.py │ ├── __pycache__ │ │ └── dota_utils.cpython-37.pyc │ ├── _polyiou.cpython-37m-x86_64-linux-gnu.so │ ├── build │ │ └── temp.linux-x86_64-3.7 │ │ │ ├── polyiou.o │ │ │ └── polyiou_wrap.o │ ├── dota_evaluation_task1.py │ ├── dota_evaluation_task2.py │ ├── dota_poly2rbox.py │ ├── dota_utils.py │ ├── hrsc2016_evaluation.py │ ├── mAOE_evaluation.py │ ├── poly_nms_gpu │ │ ├── Makefile │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ └── __init__.cpython-37.pyc │ │ ├── dist │ │ │ └── rotation-0.0.0-py3.7-linux-x86_64.egg │ │ ├── nms_wrapper.py │ │ ├── poly_nms.cpp │ │ ├── poly_nms.hpp │ │ ├── poly_nms.pyx │ │ ├── poly_nms_kernel.cu │ │ ├── poly_nms_test.py │ │ ├── poly_overlaps.cpp │ │ ├── poly_overlaps.hpp │ │ ├── poly_overlaps.pyx │ │ ├── poly_overlaps_kernel.cu │ │ ├── rotation.egg-info │ │ │ ├── PKG-INFO │ │ │ ├── SOURCES.txt │ │ │ ├── dependency_links.txt │ │ │ └── top_level.txt │ │ └── setup.py │ ├── polyiou.cpp │ ├── polyiou.h │ ├── polyiou.i │ ├── polyiou.py │ ├── polyiou_wrap.cxx │ ├── prepare_dota1_ms.py │ ├── prepare_hrsc2016.py │ ├── results_ensemble.py │ ├── results_obb2hbb.py │ ├── setup.py │ └── ucasaod_evaluation.py ├── README.md ├── data │ ├── dotav15_poly.yaml │ ├── hyps │ │ └── obb │ │ │ ├── hyp.finetune_dota.yaml │ │ │ └── hyp.finetune_dota_CloseAug.yaml │ ├── scripts │ │ └── download_weights.sh │ ├── yolov5obb_demo.yaml │ └── yolov5obb_demo_split.yaml ├── dataset │ ├── dataset_demo │ │ ├── images │ │ │ └── P0032.png │ │ ├── imgnamefile.txt │ │ └── labelTxt │ │ │ └── P0032.txt │ └── dataset_demosplit │ │ ├── images │ │ ├── P0032__1__0___0.png │ │ ├── P0032__1__0___665.png │ │ ├── P0032__1__1401___0.png │ │ ├── P0032__1__1401___665.png │ │ ├── P0032__1__924___0.png │ │ └── P0032__1__924___665.png │ │ └── labelTxt │ │ ├── P0032__1__0___0.txt │ │ ├── P0032__1__0___665.txt │ │ ├── P0032__1__1401___0.txt │ │ ├── P0032__1__1401___665.txt │ │ ├── P0032__1__924___0.txt │ │ └── P0032__1__924___665.txt ├── detect.py ├── docs │ ├── ChangeLog.md │ ├── GetStart.md │ ├── YOLOv5_README.md │ ├── detection.png │ ├── install.md │ ├── results.png │ └── train_batch6.jpg ├── export.py ├── gen_wts.py ├── hubconf.py ├── models │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── common.cpython-37.pyc │ │ ├── common.cpython-39.pyc │ │ ├── experimental.cpython-37.pyc │ │ ├── experimental.cpython-39.pyc │ │ ├── yolo.cpython-37.pyc │ │ └── yolo.cpython-39.pyc │ ├── common.py │ ├── experimental.py │ ├── hub │ │ ├── anchors.yaml │ │ ├── yolov3-spp.yaml │ │ ├── yolov3-tiny.yaml │ │ ├── yolov3.yaml │ │ ├── yolov5-bifpn.yaml │ │ ├── yolov5-fpn.yaml │ │ ├── yolov5-p2.yaml │ │ ├── yolov5-p6.yaml │ │ ├── yolov5-p7.yaml │ │ ├── yolov5-panet.yaml │ │ ├── yolov5l6.yaml │ │ ├── yolov5m6.yaml │ │ ├── yolov5n6.yaml │ │ ├── yolov5s-ghost.yaml │ │ ├── yolov5s-transformer.yaml │ │ ├── yolov5s6.yaml │ │ └── yolov5x6.yaml │ ├── tf.py │ ├── yolo.py │ ├── yolov5l.yaml │ ├── yolov5m.yaml │ ├── yolov5n.yaml │ ├── yolov5s.yaml │ └── yolov5x.yaml ├── requirements.txt ├── split_txt.py ├── tools │ └── TestJson2VocClassTxt.py ├── train.py ├── utils │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── activations.cpython-37.pyc │ │ ├── augmentations.cpython-37.pyc │ │ ├── autoanchor.cpython-37.pyc │ │ ├── autobatch.cpython-37.pyc │ │ ├── callbacks.cpython-37.pyc │ │ ├── datasets.cpython-37.pyc │ │ ├── downloads.cpython-37.pyc │ │ ├── general.cpython-37.pyc │ │ ├── loss.cpython-37.pyc │ │ ├── metrics.cpython-37.pyc │ │ ├── plots.cpython-37.pyc │ │ ├── rboxs_utils.cpython-37.pyc │ │ └── torch_utils.cpython-37.pyc │ ├── activations.py │ ├── augmentations.py │ ├── autoanchor.py │ ├── autobatch.py │ ├── aws │ │ ├── __init__.py │ │ ├── mime.sh │ │ ├── resume.py │ │ └── userdata.sh │ ├── callbacks.py │ ├── datasets.py │ ├── downloads.py │ ├── flask_rest_api │ │ ├── README.md │ │ ├── example_request.py │ │ └── restapi.py │ ├── general.py │ ├── google_app_engine │ │ ├── Dockerfile │ │ ├── additional_requirements.txt │ │ └── app.yaml │ ├── loggers │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ └── __init__.cpython-39.pyc │ │ └── wandb │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── wandb_utils.cpython-37.pyc │ │ │ └── wandb_utils.cpython-39.pyc │ │ │ ├── log_dataset.py │ │ │ ├── sweep.py │ │ │ ├── sweep.yaml │ │ │ └── wandb_utils.py │ ├── loss.py │ ├── metrics.py │ ├── nms_rotated │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-37.pyc │ │ │ └── nms_rotated_wrapper.cpython-37.pyc │ │ ├── build │ │ │ ├── lib.linux-x86_64-3.7 │ │ │ │ └── nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so │ │ │ └── temp.linux-x86_64-3.7 │ │ │ │ └── src │ │ │ │ ├── nms_rotated_cpu.o │ │ │ │ ├── nms_rotated_cuda.o │ │ │ │ ├── nms_rotated_ext.o │ │ │ │ └── poly_nms_cuda.o │ │ ├── nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so │ │ ├── nms_rotated_ext.egg-info │ │ │ ├── PKG-INFO │ │ │ ├── SOURCES.txt │ │ │ ├── dependency_links.txt │ │ │ ├── not-zip-safe │ │ │ └── top_level.txt │ │ ├── nms_rotated_wrapper.py │ │ ├── setup.py │ │ └── src │ │ │ ├── box_iou_rotated_utils.h │ │ │ ├── nms_rotated_cpu.cpp │ │ │ ├── nms_rotated_cuda.cu │ │ │ ├── nms_rotated_ext.cpp │ │ │ ├── poly_nms_cpu.cpp │ │ │ └── poly_nms_cuda.cu │ ├── plots.py │ ├── rboxs_utils.py │ └── torch_utils.py ├── val.py ├── xml2txt.py ├── yolov5s.pt └── 注意,根目录这个yolov5s.pt是预训练模型,不是最后旋转框的模型 └── yolov5_obb_tensorrt_cpp ├── CMakeLists.txt ├── README.md ├── common.hpp ├── cuda_utils.h ├── gen_wts.py ├── images └── test.jpg ├── logging.h ├── macros.h ├── result.jpg ├── yololayer.cu ├── yololayer.h ├── yolov5_gen.cpp └── yolov5_use.cpp /README.md: -------------------------------------------------------------------------------- 1 | # yolov5_for_oriented_object_detection 2 | 3 | The Pytorch implementation is [yolov5_obb](https://github.com/hukaixuan19970627/yolov5_obb ) 4 | 5 | actually ,This repo is based on [yolov5](https://github.com/ultralytics/yolov5/tree/v6.0/) and only used at Linux PC (partial dependency install hardly at windows PC) 6 | 7 | ## How to Run, yolov5s as example 8 | 1.train your datasets in {yolov5_obb} 9 | generate .pt at {yolov5_obb/runs/exp[?]/weights/best.pt} 10 | 11 | 2. generate .wts from {yolov5_obb/runs/exp[?]/weights/best.pt} 12 | 13 | ``` 14 | cp gen_wts.py {yolov5_obb} 15 | python gen_wts.py -w runs/train/exp[]/weights/best.pt -o yolov5s.wts 16 | // a file 'yolov5s.wts' will be generated. 17 | ``` 18 | 19 | 3. build **tensorrt engine ** 20 | 21 | ``` 22 | cd {yolov5_obb_tensorrt_cpp} 23 | // update CLASS_NUM/Inputsize in yololayer.h 24 | // caution:CLASS_NUM= your classes +180(180 for angle classes) 25 | mkdir build 26 | cd build 27 | cp {yolov5_obb}/yolov5s.wts ../yolov5s.wts 28 | cmake .. 29 | make 30 | sudo ./yolov5_gen -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file 31 | // For example yolov5s 32 | sudo ./yolov5_gen -s ../yolov5s.wts ../yolov5s.engine s 33 | ``` 34 | 35 | 4. use **tensorr engine** 36 | ``` 37 | sudo ./yolov5_use ../yolov5s.engine ../images/test.jpg 38 | ``` 39 | ![image](https://github.com/fish-kong/Yolov5-obb-Tensorrt-Infer/blob/main/yolov5_obb_tensorrt_cpp/result.jpg) 40 | 41 | -------------------------------------------------------------------------------- /yolov5_obb/Arial.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/Arial.ttf -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/DOTA.py: -------------------------------------------------------------------------------- 1 | #The code is used for visulization, inspired from cocoapi 2 | # Licensed under the Simplified BSD License [see bsd.txt] 3 | 4 | import os 5 | import matplotlib.pyplot as plt 6 | from matplotlib.collections import PatchCollection 7 | from matplotlib.patches import Polygon, Circle 8 | import numpy as np 9 | import dota_utils as util 10 | from collections import defaultdict 11 | import cv2 12 | 13 | def _isArrayLike(obj): 14 | if type(obj) == str: 15 | return False 16 | return hasattr(obj, '__iter__') and hasattr(obj, '__len__') 17 | 18 | class DOTA: 19 | def __init__(self, basepath): 20 | self.basepath = basepath 21 | self.labelpath = os.path.join(basepath, 'labelTxt') 22 | self.imagepath = os.path.join(basepath, 'images') 23 | self.imgpaths = util.GetFileFromThisRootDir(self.labelpath) 24 | self.imglist = [util.custombasename(x) for x in self.imgpaths] 25 | self.catToImgs = defaultdict(list) 26 | self.ImgToAnns = defaultdict(list) 27 | self.createIndex() 28 | 29 | def createIndex(self): 30 | for filename in self.imgpaths: 31 | objects = util.parse_dota_poly(filename) 32 | imgid = util.custombasename(filename) 33 | self.ImgToAnns[imgid] = objects 34 | for obj in objects: 35 | cat = obj['name'] 36 | self.catToImgs[cat].append(imgid) 37 | 38 | def getImgIds(self, catNms=[]): 39 | """ 40 | :param catNms: category names 41 | :return: all the image ids contain the categories 42 | """ 43 | catNms = catNms if _isArrayLike(catNms) else [catNms] 44 | if len(catNms) == 0: 45 | return self.imglist 46 | else: 47 | imgids = [] 48 | for i, cat in enumerate(catNms): 49 | if i == 0: 50 | imgids = set(self.catToImgs[cat]) 51 | else: 52 | imgids &= set(self.catToImgs[cat]) 53 | return list(imgids) 54 | 55 | def loadAnns(self, catNms=[], imgId = None, difficult=None): 56 | """ 57 | :param catNms: category names 58 | :param imgId: the img to load anns 59 | :return: objects 60 | """ 61 | catNms = catNms if _isArrayLike(catNms) else [catNms] 62 | objects = self.ImgToAnns[imgId] 63 | if len(catNms) == 0: 64 | return objects 65 | outobjects = [obj for obj in objects if (obj['name'] in catNms)] 66 | return outobjects 67 | def showAnns(self, objects, imgId, range): 68 | """ 69 | :param catNms: category names 70 | :param objects: objects to show 71 | :param imgId: img to show 72 | :param range: display range in the img 73 | :return: 74 | """ 75 | img = self.loadImgs(imgId)[0] 76 | plt.imshow(img) 77 | plt.axis('off') 78 | 79 | ax = plt.gca() 80 | ax.set_autoscale_on(False) 81 | polygons = [] 82 | color = [] 83 | circles = [] 84 | r = 5 85 | for obj in objects: 86 | c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] 87 | poly = obj['poly'] 88 | polygons.append(Polygon(poly)) 89 | color.append(c) 90 | point = poly[0] 91 | circle = Circle((point[0], point[1]), r) 92 | circles.append(circle) 93 | p = PatchCollection(polygons, facecolors=color, linewidths=0, alpha=0.4) 94 | ax.add_collection(p) 95 | p = PatchCollection(polygons, facecolors='none', edgecolors=color, linewidths=2) 96 | ax.add_collection(p) 97 | p = PatchCollection(circles, facecolors='red') 98 | ax.add_collection(p) 99 | def loadImgs(self, imgids=[]): 100 | """ 101 | :param imgids: integer ids specifying img 102 | :return: loaded img objects 103 | """ 104 | print('isarralike:', _isArrayLike(imgids)) 105 | imgids = imgids if _isArrayLike(imgids) else [imgids] 106 | print('imgids:', imgids) 107 | imgs = [] 108 | for imgid in imgids: 109 | filename = os.path.join(self.imagepath, imgid + '.png') 110 | print('filename:', filename) 111 | img = cv2.imread(filename) 112 | imgs.append(img) 113 | return imgs 114 | 115 | if __name__ == '__main__': 116 | examplesplit = DOTA('/media/test/4d846cae-2315-4928-8d1b-ca6d3a61a3c6/DOTA/DOTAv2.0/trainval_split_1024') 117 | imgids = examplesplit.getImgIds(catNms=['plane']) 118 | img = examplesplit.loadImgs(imgids) 119 | for imgid in imgids: 120 | anns = examplesplit.loadAnns(imgId=imgid) 121 | examplesplit.showAnns(anns, imgid, 2) -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/DOTA2COCO.py: -------------------------------------------------------------------------------- 1 | import dota_utils as util 2 | import os 3 | import cv2 4 | import json 5 | from PIL import Image 6 | import os.path as osp 7 | 8 | wordname_15 = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 'tennis-court', 9 | 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter'] 10 | 11 | wordname_16 = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 'tennis-court', 12 | 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter', 'container-crane'] 13 | 14 | wordname_18 = [ 'plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 15 | 'tennis-court', 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 16 | 'swimming-pool', 'helicopter', 'container-crane', 'airport', 'helipad'] 17 | 18 | hrsc_2016 = ['ship'] 19 | ucas_aod = ['car', 'airplane'] 20 | 21 | def DOTA2COCOTrain(srcpath, destfile, cls_names, difficult='2'): 22 | # set difficult to filter '2', '1', or do not filter, set '-1' 23 | 24 | imageparent = os.path.join(srcpath, 'images') 25 | labelparent = os.path.join(srcpath, 'labelTxt') 26 | 27 | data_dict = {} 28 | data_dict['images'] = [] 29 | data_dict['categories'] = [] 30 | data_dict['annotations'] = [] 31 | for idex, name in enumerate(cls_names): 32 | single_cat = {'id': idex + 1, 'name': name, 'supercategory': name} 33 | data_dict['categories'].append(single_cat) 34 | 35 | inst_count = 1 36 | image_id = 1 37 | with open(destfile, 'w') as f_out: 38 | filenames = util.GetFileFromThisRootDir(labelparent) 39 | for file in filenames: 40 | basename = util.custombasename(file) 41 | # image_id = int(basename[1:]) 42 | 43 | imagepath = os.path.join(imageparent, basename + '.png') 44 | img = cv2.imread(imagepath) 45 | height, width, c = img.shape 46 | 47 | single_image = {} 48 | single_image['file_name'] = basename + '.png' 49 | single_image['id'] = image_id 50 | single_image['width'] = width 51 | single_image['height'] = height 52 | data_dict['images'].append(single_image) 53 | 54 | # annotations 55 | objects = util.parse_dota_poly2(file) 56 | for obj in objects: 57 | if obj['difficult'] == difficult: 58 | print('difficult: ', difficult) 59 | continue 60 | single_obj = {} 61 | single_obj['area'] = obj['area'] 62 | 63 | if obj['name'] not in cls_names: # 去掉别的类 64 | print('This classname not found: ', obj['name']) 65 | continue 66 | 67 | single_obj['category_id'] = cls_names.index(obj['name']) + 1 68 | single_obj['segmentation'] = [] 69 | single_obj['segmentation'].append(obj['poly']) 70 | single_obj['iscrowd'] = 0 71 | # xmin, ymin, xmax, ymax = min(obj['poly'][0::2]), min(obj['poly'][1::2]), \ 72 | # max(obj['poly'][0::2]), max(obj['poly'][1::2]) 73 | x1 = obj['poly'][0] 74 | y1 = obj['poly'][1] 75 | x2 = obj['poly'][2] 76 | y2 = obj['poly'][3] 77 | x3 = obj['poly'][4] 78 | y3 = obj['poly'][5] 79 | x4 = obj['poly'][6] 80 | y4 = obj['poly'][7] 81 | 82 | # width, height = xmax - xmin, ymax - ymin 83 | single_obj['bbox'] = x1, y1, x2, y2, x3, y3, x4, y4 84 | single_obj['image_id'] = image_id 85 | data_dict['annotations'].append(single_obj) 86 | single_obj['id'] = inst_count 87 | inst_count = inst_count + 1 88 | image_id = image_id + 1 89 | json.dump(data_dict, f_out) 90 | 91 | def DOTA2COCOTest(srcpath, destfile, cls_names): 92 | imageparent = os.path.join(srcpath, 'images') 93 | data_dict = {} 94 | 95 | data_dict['images'] = [] 96 | data_dict['categories'] = [] 97 | for idex, name in enumerate(cls_names): 98 | single_cat = {'id': idex + 1, 'name': name, 'supercategory': name} 99 | data_dict['categories'].append(single_cat) 100 | 101 | image_id = 1 102 | with open(destfile, 'w') as f_out: 103 | filenames = util.GetFileFromThisRootDir(imageparent) 104 | for file in filenames: 105 | basename = util.custombasename(file) 106 | imagepath = os.path.join(imageparent, basename + '.png') 107 | img = Image.open(imagepath) 108 | height = img.height 109 | width = img.width 110 | 111 | single_image = {} 112 | single_image['file_name'] = basename + '.png' 113 | single_image['id'] = image_id 114 | single_image['width'] = width 115 | single_image['height'] = height 116 | data_dict['images'].append(single_image) 117 | 118 | image_id = image_id + 1 119 | json.dump(data_dict, f_out) 120 | 121 | if __name__ == '__main__': 122 | 123 | DOTA2COCOTrain(r'data/dataset_demo_split', 124 | osp.join('data/dataset_demo_split', 'train_datasetdemo.json'), 125 | wordname_15) 126 | print('Train DONE') 127 | # DOTA2COCOTrain(r'/dataset/Dota/Dota_V1.0/train_sp/', 128 | # r'/mnt/SSD/lwt_workdir/s2anet_branch/data/ucas_aod/Train/trainval_coco_8point.json', 129 | # ucas_aod) 130 | DOTA2COCOTest(r'data/dataset_demo_split', 131 | osp.join('data/dataset_demo_split', 'test_datasetdemo.json'), 132 | wordname_15) 133 | # DOTA2COCOTest(r'/mnt/SSD/lwt_workdir/s2anet_branch/data/ucas_aod/Test/', 134 | # r'/mnt/SSD/lwt_workdir/s2anet_branch/data/ucas_aod/Test/test_coco_8point.json', 135 | # ucas_aod) 136 | 137 | print('Test DONE') 138 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/DOTA2JSON.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | import os.path as osp 4 | import random 5 | 6 | from PIL import Image 7 | 8 | from dota_poly2rbox import poly2rbox_single_v2, poly2rbox_single 9 | 10 | 11 | def parse_ann_info(img_base_path, label_base_path, img_name): 12 | lab_path = osp.join(label_base_path, img_name+'.txt') 13 | bboxes, labels, bboxes_ignore, labels_ignore = [], [], [], [] 14 | with open(lab_path, 'r') as f: 15 | for ann_line in f.readlines(): 16 | ann_line = ann_line.strip().split(' ') 17 | bbox = [float(ann_line[i]) for i in range(8)] 18 | # 8 point to 5 point xywha 19 | bbox = poly2rbox_single_v2(bbox) 20 | class_name = ann_line[8] 21 | difficult = int(ann_line[9]) 22 | # ignore difficult =2 23 | if difficult == 0: 24 | bboxes.append(bbox) 25 | labels.append(class_name) 26 | elif difficult == 1: 27 | bboxes_ignore.append(bbox) 28 | labels_ignore.append(class_name) 29 | return bboxes, labels, bboxes_ignore, labels_ignore 30 | 31 | 32 | def generate_txt_labels(src_path, out_path, trainval=True): 33 | """Generate .txt labels recording img_names 34 | Args: 35 | src_path: dataset path containing images and labelTxt folders. 36 | out_path: output txt file path 37 | trainval: trainval or test? 38 | """ 39 | img_path = os.path.join(src_path, 'images') 40 | label_path = os.path.join(src_path, 'labelTxt') 41 | img_lists = os.listdir(img_path) 42 | with open(out_path, 'w') as f: 43 | for img in img_lists: 44 | img_name = osp.splitext(img)[0] 45 | label = os.path.join(label_path, img_name+'.txt') 46 | if(trainval == True): 47 | if(os.path.exists(label) == False): 48 | print('Label:'+img_name+'.txt'+' Not Exist') 49 | else: 50 | f.write(img_name+'\n') 51 | else: 52 | f.write(img_name+'\n') 53 | 54 | 55 | def generate_json_labels(src_path, out_path, trainval=True): 56 | """Generate .json labels which is similar to coco format 57 | Args: 58 | src_path: dataset path containing images and labelTxt folders. 59 | out_path: output json file path 60 | trainval: trainval or test? 61 | """ 62 | img_path = os.path.join(src_path, 'images') 63 | label_path = os.path.join(src_path, 'labelTxt') 64 | img_lists = os.listdir(img_path) 65 | 66 | data_dict = [] 67 | 68 | with open(out_path, 'w') as f: 69 | for id, img in enumerate(img_lists): 70 | img_info = {} 71 | img_name = osp.splitext(img)[0] 72 | label = os.path.join(label_path, img_name+'.txt') 73 | img = Image.open(osp.join(img_path, img)) 74 | img_info['filename'] = img_name+'.png' 75 | img_info['height'] = img.height 76 | img_info['width'] = img.width 77 | img_info['id'] = id 78 | if(trainval == True): 79 | if(os.path.exists(label) == False): 80 | print('Label:'+img_name+'.txt'+' Not Exist') 81 | else: 82 | bboxes, labels, bboxes_ignore, labels_ignore = parse_ann_info( 83 | img_path, label_path, img_name) 84 | ann = {} 85 | ann['bboxes'] = bboxes 86 | ann['labels'] = labels 87 | ann['bboxes_ignore'] = bboxes_ignore 88 | ann['labels_ignore'] = labels_ignore 89 | img_info['annotations'] = ann 90 | data_dict.append(img_info) 91 | json.dump(data_dict, f) 92 | 93 | 94 | if __name__ == '__main__': 95 | generate_json_labels('/data1/dataset_demo/DOTA_demo/', 96 | '/data1/OrientedRepPoints/data/dota/trainval_split/trainval.json') 97 | generate_json_labels('/data1/dataset_demo/DOTA_demo/', 98 | '/data1/OrientedRepPoints/data/dota/test_split/test.json', trainval=False) 99 | print('done!') 100 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/ResultMerge.py: -------------------------------------------------------------------------------- 1 | """ 2 | To use the code, users should to config detpath, annopath and imagesetfile 3 | detpath is the path for 15 result files, for the format, you can refer to "http://captain.whu.edu.cn/DOTAweb/tasks.html" 4 | search for PATH_TO_BE_CONFIGURED to config the paths 5 | Note, the evaluation is on the large scale images 6 | """ 7 | import os 8 | import numpy as np 9 | import dota_utils as util 10 | import re 11 | import time 12 | import polyiou 13 | 14 | # the thresh for nms when merge image 15 | nms_thresh = 0.3 16 | 17 | 18 | def py_cpu_nms_poly(dets, thresh): 19 | scores = dets[:, 8] 20 | polys = [] 21 | areas = [] 22 | for i in range(len(dets)): 23 | tm_polygon = polyiou.VectorDouble([dets[i][0], dets[i][1], 24 | dets[i][2], dets[i][3], 25 | dets[i][4], dets[i][5], 26 | dets[i][6], dets[i][7]]) 27 | polys.append(tm_polygon) 28 | order = scores.argsort()[::-1] 29 | 30 | keep = [] 31 | while order.size > 0: 32 | ovr = [] 33 | i = order[0] 34 | keep.append(i) 35 | for j in range(order.size - 1): 36 | iou = polyiou.iou_poly(polys[i], polys[order[j + 1]]) 37 | ovr.append(iou) 38 | ovr = np.array(ovr) 39 | inds = np.where(ovr <= thresh)[0] 40 | order = order[inds + 1] 41 | return keep 42 | 43 | 44 | def py_cpu_nms(dets, thresh): 45 | """Pure Python NMS baseline.""" 46 | #print('dets:', dets) 47 | x1 = dets[:, 0] 48 | y1 = dets[:, 1] 49 | x2 = dets[:, 2] 50 | y2 = dets[:, 3] 51 | scores = dets[:, 4] 52 | 53 | areas = (x2 - x1 + 1) * (y2 - y1 + 1) 54 | # index for dets 55 | order = scores.argsort()[::-1] 56 | 57 | keep = [] 58 | while order.size > 0: 59 | i = order[0] 60 | keep.append(i) 61 | xx1 = np.maximum(x1[i], x1[order[1:]]) 62 | yy1 = np.maximum(y1[i], y1[order[1:]]) 63 | xx2 = np.minimum(x2[i], x2[order[1:]]) 64 | yy2 = np.minimum(y2[i], y2[order[1:]]) 65 | 66 | w = np.maximum(0.0, xx2 - xx1 + 1) 67 | h = np.maximum(0.0, yy2 - yy1 + 1) 68 | inter = w * h 69 | ovr = inter / (areas[i] + areas[order[1:]] - inter) 70 | 71 | inds = np.where(ovr <= thresh)[0] 72 | order = order[inds + 1] 73 | 74 | return keep 75 | 76 | 77 | def nmsbynamedict(nameboxdict, nms, thresh): 78 | nameboxnmsdict = {x: [] for x in nameboxdict} 79 | for imgname in nameboxdict: 80 | #print('imgname:', imgname) 81 | #keep = py_cpu_nms(np.array(nameboxdict[imgname]), thresh) 82 | #print('type nameboxdict:', type(nameboxnmsdict)) 83 | #print('type imgname:', type(imgname)) 84 | #print('type nms:', type(nms)) 85 | keep = nms(np.array(nameboxdict[imgname]), thresh) 86 | #print('keep:', keep) 87 | outdets = [] 88 | #print('nameboxdict[imgname]: ', nameboxnmsdict[imgname]) 89 | for index in keep: 90 | # print('index:', index) 91 | outdets.append(nameboxdict[imgname][index]) 92 | nameboxnmsdict[imgname] = outdets 93 | return nameboxnmsdict 94 | 95 | 96 | def poly2origpoly(poly, x, y, rate): 97 | origpoly = [] 98 | for i in range(int(len(poly)/2)): 99 | tmp_x = float(poly[i * 2] + x) / float(rate) 100 | tmp_y = float(poly[i * 2 + 1] + y) / float(rate) 101 | origpoly.append(tmp_x) 102 | origpoly.append(tmp_y) 103 | return origpoly 104 | 105 | 106 | def mergebase(srcpath, dstpath, nms): 107 | filelist = util.GetFileFromThisRootDir(srcpath) 108 | for fullname in filelist: 109 | name = util.custombasename(fullname) 110 | #print('name:', name) 111 | dstname = os.path.join(dstpath, name + '.txt') 112 | with open(fullname, 'r') as f_in: 113 | nameboxdict = {} 114 | lines = f_in.readlines() 115 | splitlines = [x.strip().split(' ') for x in lines] 116 | for splitline in splitlines: 117 | subname = splitline[0] 118 | splitname = subname.split('__') 119 | oriname = splitname[0] 120 | pattern1 = re.compile(r'__\d+___\d+') 121 | #print('subname:', subname) 122 | x_y = re.findall(pattern1, subname) 123 | x_y_2 = re.findall(r'\d+', x_y[0]) 124 | x, y = int(x_y_2[0]), int(x_y_2[1]) 125 | 126 | pattern2 = re.compile(r'__([\d+\.]+)__\d+___') 127 | 128 | rate = re.findall(pattern2, subname)[0] 129 | 130 | confidence = splitline[1] 131 | poly = list(map(float, splitline[2:])) 132 | origpoly = poly2origpoly(poly, x, y, rate) 133 | det = origpoly 134 | det.append(confidence) 135 | det = list(map(float, det)) 136 | if (oriname not in nameboxdict): 137 | nameboxdict[oriname] = [] 138 | nameboxdict[oriname].append(det) 139 | nameboxnmsdict = nmsbynamedict(nameboxdict, nms, nms_thresh) 140 | with open(dstname, 'w') as f_out: 141 | for imgname in nameboxnmsdict: 142 | for det in nameboxnmsdict[imgname]: 143 | #print('det:', det) 144 | confidence = det[-1] 145 | bbox = det[0:-1] 146 | outline = imgname + ' ' + \ 147 | str(confidence) + ' ' + ' '.join(map(str, bbox)) 148 | #print('outline:', outline) 149 | f_out.write(outline + '\n') 150 | 151 | 152 | def mergebyrec(srcpath, dstpath): 153 | """ 154 | srcpath: result files before merge and nms 155 | dstpath: result files after merge and nms 156 | """ 157 | # srcpath = r'E:\bod-dataset\results\bod-v3_rfcn_2000000' 158 | # dstpath = r'E:\bod-dataset\results\bod-v3_rfcn_2000000_nms' 159 | 160 | mergebase(srcpath, 161 | dstpath, 162 | py_cpu_nms) 163 | 164 | 165 | def mergebypoly(srcpath, dstpath): 166 | """ 167 | srcpath: result files before merge and nms 168 | dstpath: result files after merge and nms 169 | """ 170 | mergebase(srcpath, 171 | dstpath, 172 | py_cpu_nms_poly) 173 | 174 | 175 | if __name__ == '__main__': 176 | # see demo for example 177 | mergebypoly(r'work_dirs/temp/result_raw',\ 178 | r'work_dirs/temp/result_merge') 179 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/SplitOnlyImage.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | import copy 5 | import dota_utils as util 6 | 7 | class splitbase(): 8 | def __init__(self, 9 | srcpath, 10 | dstpath, 11 | gap=100, 12 | subsize=1024, 13 | ext='.png'): 14 | self.srcpath = srcpath 15 | self.outpath = dstpath 16 | self.gap = gap 17 | self.subsize = subsize 18 | self.slide = self.subsize - self.gap 19 | self.srcpath = srcpath 20 | self.dstpath = dstpath 21 | self.ext = ext 22 | def saveimagepatches(self, img, subimgname, left, up, ext='.png'): 23 | subimg = copy.deepcopy(img[up: (up + self.subsize), left: (left + self.subsize)]) 24 | outdir = os.path.join(self.dstpath, subimgname + ext) 25 | cv2.imwrite(outdir, subimg) 26 | 27 | def SplitSingle(self, name, rate, extent): 28 | img = cv2.imread(os.path.join(self.srcpath, name + extent)) 29 | assert np.shape(img) != () 30 | 31 | if (rate != 1): 32 | resizeimg = cv2.resize(img, None, fx=rate, fy=rate, interpolation = cv2.INTER_CUBIC) 33 | else: 34 | resizeimg = img 35 | outbasename = name + '__' + str(rate) + '__' 36 | 37 | weight = np.shape(resizeimg)[1] 38 | height = np.shape(resizeimg)[0] 39 | 40 | left, up = 0, 0 41 | while (left < weight): 42 | if (left + self.subsize >= weight): 43 | left = max(weight - self.subsize, 0) 44 | up = 0 45 | while (up < height): 46 | if (up + self.subsize >= height): 47 | up = max(height - self.subsize, 0) 48 | subimgname = outbasename + str(left) + '___' + str(up) 49 | self.saveimagepatches(resizeimg, subimgname, left, up) 50 | if (up + self.subsize >= height): 51 | break 52 | else: 53 | up = up + self.slide 54 | if (left + self.subsize >= weight): 55 | break 56 | else: 57 | left = left + self.slide 58 | 59 | def splitdata(self, rate): 60 | 61 | imagelist = util.GetFileFromThisRootDir(self.srcpath) 62 | imagenames = [util.custombasename(x) for x in imagelist if (util.custombasename(x) != 'Thumbs')] 63 | for name in imagenames: 64 | self.SplitSingle(name, rate, self.ext) 65 | if __name__ == '__main__': 66 | split = splitbase(r'example/images', 67 | r'example/imagesSplit') 68 | split.splitdata(1) -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/SplitOnlyImage_multi_process.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | import copy 5 | import dota_utils as util 6 | from multiprocessing import Pool 7 | from functools import partial 8 | 9 | 10 | def split_single_warp(name, split_base, rate, extent): 11 | split_base.SplitSingle(name, rate, extent) 12 | 13 | 14 | class splitbase(): 15 | def __init__(self, 16 | srcpath, 17 | dstpath, 18 | gap=100, 19 | subsize=1024, 20 | ext='.png', 21 | padding=True, 22 | num_process=32): 23 | self.srcpath = srcpath 24 | self.outpath = dstpath 25 | self.gap = gap 26 | self.subsize = subsize 27 | self.slide = self.subsize - self.gap 28 | self.srcpath = srcpath 29 | self.dstpath = dstpath 30 | self.ext = ext 31 | self.padding = padding 32 | self.pool = Pool(num_process) 33 | 34 | if not os.path.isdir(self.outpath): 35 | os.mkdir(self.outpath) 36 | 37 | def saveimagepatches(self, img, subimgname, left, up, ext='.png'): 38 | subimg = copy.deepcopy( 39 | img[up: (up + self.subsize), left: (left + self.subsize)]) 40 | outdir = os.path.join(self.dstpath, subimgname + ext) 41 | h, w, c = np.shape(subimg) 42 | if (self.padding): 43 | outimg = np.zeros((self.subsize, self.subsize, 3)) 44 | outimg[0:h, 0:w, :] = subimg 45 | cv2.imwrite(outdir, outimg) 46 | else: 47 | cv2.imwrite(outdir, subimg) 48 | 49 | def SplitSingle(self, name, rate, extent): 50 | img = cv2.imread(os.path.join(self.srcpath, name + extent)) 51 | assert np.shape(img) != () 52 | 53 | if (rate != 1): 54 | resizeimg = cv2.resize( 55 | img, None, fx=rate, fy=rate, interpolation=cv2.INTER_CUBIC) 56 | else: 57 | resizeimg = img 58 | outbasename = name + '__' + str(rate) + '__' 59 | 60 | weight = np.shape(resizeimg)[1] 61 | height = np.shape(resizeimg)[0] 62 | 63 | # if (max(weight, height) < self.subsize/2): 64 | # return 65 | 66 | left, up = 0, 0 67 | while (left < weight): 68 | if (left + self.subsize >= weight): 69 | left = max(weight - self.subsize, 0) 70 | up = 0 71 | while (up < height): 72 | if (up + self.subsize >= height): 73 | up = max(height - self.subsize, 0) 74 | subimgname = outbasename + str(left) + '___' + str(up) 75 | self.saveimagepatches(resizeimg, subimgname, left, up) 76 | if (up + self.subsize >= height): 77 | break 78 | else: 79 | up = up + self.slide 80 | if (left + self.subsize >= weight): 81 | break 82 | else: 83 | left = left + self.slide 84 | 85 | def splitdata(self, rate): 86 | 87 | imagelist = util.GetFileFromThisRootDir(self.srcpath) 88 | imagenames = [util.custombasename(x) for x in imagelist if ( 89 | util.custombasename(x) != 'Thumbs')] 90 | 91 | # worker = partial(self.SplitSingle, rate=rate, extent=self.ext) 92 | worker = partial(split_single_warp, split_base=self, 93 | rate=rate, extent=self.ext) 94 | self.pool.map(worker, imagenames) 95 | # 96 | # for name in imagenames: 97 | # self.SplitSingle(name, rate, self.ext) 98 | 99 | def __getstate__(self): 100 | self_dict = self.__dict__.copy() 101 | del self_dict['pool'] 102 | return self_dict 103 | 104 | def __setstate__(self, state): 105 | self.__dict__.update(state) 106 | 107 | 108 | if __name__ == '__main__': 109 | split = splitbase(r'/media/test/4d846cae-2315-4928-8d1b-ca6d3a61a3c6/DOTA/DOTAv1.5/test/images', 110 | r'/media/test/4d846cae-2315-4928-8d1b-ca6d3a61a3c6/DOTA/DOTAv2.0/test-dev_split/images', 111 | gap=200, subsize=1024, num_process=8) 112 | split.splitdata(1) 113 | # split.splitdata(0.5) 114 | # split.splitdata(1.5) 115 | print("Split Done!") 116 | 117 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/Visuallize_dota.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | path='/home/addtion_storage/huangtao/tomato_detect_rotate/split/images/0027__1__312___312.jpg' 5 | annpath=path.replace('images','labelTxt').replace('jpg','txt') 6 | img=cv2.imread(path) 7 | 8 | with open(annpath,'r') as f: 9 | ann=f.readlines() 10 | for each in ann: 11 | each=each.strip() 12 | axis=[float(i) for i in each.split(' ')[:8]] 13 | label=each.split(' ')[8] 14 | area1 = np.array([[axis[0], axis[1]], [axis[2], axis[3]], [axis[4], axis[5]], [axis[6], axis[7]]], dtype=np.int32) 15 | #cv2.fillPoly(img, [area1], (2, 255, 255)) 16 | cv2.polylines(img, [area1], True, (98, 9, 11), 3) 17 | plt.imshow(img[:,:,::-1]) 18 | plt.show() -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/__init__.py -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/__pycache__/dota_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/__pycache__/dota_utils.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/_polyiou.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/_polyiou.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/build/temp.linux-x86_64-3.7/polyiou.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/build/temp.linux-x86_64-3.7/polyiou.o -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/build/temp.linux-x86_64-3.7/polyiou_wrap.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/build/temp.linux-x86_64-3.7/polyiou_wrap.o -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/Makefile: -------------------------------------------------------------------------------- 1 | all: 2 | python setup.py build_ext --inplace 3 | rm -rf build 4 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/__init__.py: -------------------------------------------------------------------------------- 1 | from .poly_overlaps import poly_overlaps 2 | 3 | __all__ = ['poly_overlaps', 'poly_nms'] 4 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/poly_nms_gpu/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/dist/rotation-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/poly_nms_gpu/dist/rotation-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/nms_wrapper.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | # from nms.gpu_nms import gpu_nms 9 | # from nms.cpu_nms import cpu_nms 10 | from .poly_nms import poly_gpu_nms 11 | def poly_nms_gpu(dets, thresh, force_cpu=False): 12 | """Dispatch to either CPU or GPU NMS implementations.""" 13 | 14 | if dets.shape[0] == 0: 15 | return [] 16 | return poly_gpu_nms(dets, thresh, device_id=0) 17 | 18 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/poly_nms.hpp: -------------------------------------------------------------------------------- 1 | // 2 | // Created by dingjian on 18-5-24. 3 | // 4 | 5 | #ifndef DOTA_DEVKIT_POLY_NMS_HPP 6 | #define DOTA_DEVKIT_POLY_NMS_HPP 7 | 8 | 9 | void _poly_nms(int* keep_out, int* num_out, const float* polys_host, int polys_num, 10 | int polys_dim, float nms_overlap_thresh, int device_id); 11 | 12 | #endif //DOTA_DEVKIT_POLY_NMS_HPP 13 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/poly_nms.pyx: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | cimport numpy as np 3 | 4 | assert sizeof(int) == sizeof(np.int32_t) 5 | 6 | cdef extern from "poly_nms.hpp": 7 | void _poly_nms(np.int32_t*, int*, np.float32_t*, int, int, float, int) 8 | 9 | def poly_gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, 10 | np.int32_t device_id=0): 11 | cdef int boxes_num = dets.shape[0] 12 | cdef int boxes_dim = dets.shape[1] 13 | cdef int num_out 14 | cdef np.ndarray[np.int32_t, ndim=1] \ 15 | keep = np.zeros(boxes_num, dtype=np.int32) 16 | cdef np.ndarray[np.float32_t, ndim=1] \ 17 | scores = dets[:, 8] 18 | cdef np.ndarray[np.int_t, ndim=1] \ 19 | order = scores.argsort()[::-1] 20 | cdef np.ndarray[np.float32_t, ndim=2] \ 21 | sorted_dets = dets[order, :] 22 | _poly_nms(&keep[0], &num_out, &sorted_dets[0, 0], boxes_num, boxes_dim, thresh, device_id) 23 | keep = keep[:num_out] 24 | return list(order[keep]) 25 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/poly_nms_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/DOTA_devkit/poly_nms_gpu/poly_nms_test.py -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/poly_overlaps.hpp: -------------------------------------------------------------------------------- 1 | void _overlaps(float* overlaps,const float* boxes,const float* query_boxes, int n, int k, int device_id); 2 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/poly_overlaps.pyx: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | cimport numpy as np 3 | 4 | cdef extern from "poly_overlaps.hpp": 5 | void _overlaps(np.float32_t*, np.float32_t*, np.float32_t*, int, int, int) 6 | 7 | def poly_overlaps (np.ndarray[np.float32_t, ndim=2] boxes, np.ndarray[np.float32_t, ndim=2] query_boxes, np.int32_t device_id=0): 8 | cdef int N = boxes.shape[0] 9 | cdef int K = query_boxes.shape[0] 10 | cdef np.ndarray[np.float32_t, ndim=2] overlaps = np.zeros((N, K), dtype = np.float32) 11 | _overlaps(&overlaps[0, 0], &boxes[0, 0], &query_boxes[0, 0], N, K, device_id) 12 | return overlaps 13 | 14 | 15 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/rotation.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: rotation 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/rotation.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | poly_nms.pyx 2 | poly_nms_kernel.cu 3 | poly_overlaps.pyx 4 | poly_overlaps_kernel.cu 5 | setup.py 6 | rotation.egg-info/PKG-INFO 7 | rotation.egg-info/SOURCES.txt 8 | rotation.egg-info/dependency_links.txt 9 | rotation.egg-info/top_level.txt -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/rotation.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/rotation.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | poly_nms 2 | poly_overlaps 3 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/poly_nms_gpu/setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | setup.py file for SWIG example 3 | """ 4 | import os 5 | from os.path import join as pjoin 6 | from setuptools import setup 7 | from distutils.extension import Extension 8 | from Cython.Distutils import build_ext 9 | import subprocess 10 | import numpy as np 11 | 12 | def find_in_path(name, path): 13 | "Find a file in a search path" 14 | # Adapted fom 15 | # http://code.activestate.com/recipes/52224-find-a-file-given-a-search-path/ 16 | for dir in path.split(os.pathsep): 17 | binpath = pjoin(dir, name) 18 | if os.path.exists(binpath): 19 | return os.path.abspath(binpath) 20 | return None 21 | 22 | 23 | def locate_cuda(): 24 | """Locate the CUDA environment on the system 25 | 26 | Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64' 27 | and values giving the absolute path to each directory. 28 | 29 | Starts by looking for the CUDAHOME env variable. If not found, everything 30 | is based on finding 'nvcc' in the PATH. 31 | """ 32 | 33 | # first check if the CUDAHOME env variable is in use 34 | if 'CUDAHOME' in os.environ: 35 | home = os.environ['CUDAHOME'] 36 | nvcc = pjoin(home, 'bin', 'nvcc') 37 | else: 38 | # otherwise, search the PATH for NVCC 39 | default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin') 40 | nvcc = find_in_path('nvcc', os.environ['PATH'] + os.pathsep + default_path) 41 | if nvcc is None: 42 | raise EnvironmentError('The nvcc binary could not be ' 43 | 'located in your $PATH. Either add it to your path, or set $CUDAHOME') 44 | home = os.path.dirname(os.path.dirname(nvcc)) 45 | 46 | cudaconfig = {'home':home, 'nvcc':nvcc, 47 | 'include': pjoin(home, 'include'), 48 | 'lib64': pjoin(home, 'lib64')} 49 | try: 50 | for k, v in cudaconfig.iteritems(): 51 | if not os.path.exists(v): 52 | raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v)) 53 | except: 54 | for k, v in cudaconfig.items(): 55 | if not os.path.exists(v): 56 | raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v)) 57 | return cudaconfig 58 | CUDA = locate_cuda() 59 | 60 | 61 | # Obtain the numpy include directory. This logic works across numpy versions. 62 | try: 63 | numpy_include = np.get_include() 64 | except AttributeError: 65 | numpy_include = np.get_numpy_include() 66 | 67 | def customize_compiler_for_nvcc(self): 68 | """inject deep into distutils to customize how the dispatch 69 | to gcc/nvcc works. 70 | 71 | If you subclass UnixCCompiler, it's not trivial to get your subclass 72 | injected in, and still have the right customizations (i.e. 73 | distutils.sysconfig.customize_compiler) run on it. So instead of going 74 | the OO route, I have this. Note, it's kindof like a wierd functional 75 | subclassing going on.""" 76 | 77 | # tell the compiler it can processes .cu 78 | self.src_extensions.append('.cu') 79 | 80 | # save references to the default compiler_so and _comple methods 81 | default_compiler_so = self.compiler_so 82 | super = self._compile 83 | 84 | # now redefine the _compile method. This gets executed for each 85 | # object but distutils doesn't have the ability to change compilers 86 | # based on source extension: we add it. 87 | def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts): 88 | if os.path.splitext(src)[1] == '.cu': 89 | # use the cuda for .cu files 90 | self.set_executable('compiler_so', CUDA['nvcc']) 91 | # use only a subset of the extra_postargs, which are 1-1 translated 92 | # from the extra_compile_args in the Extension class 93 | postargs = extra_postargs['nvcc'] 94 | else: 95 | postargs = extra_postargs['gcc'] 96 | 97 | super(obj, src, ext, cc_args, postargs, pp_opts) 98 | # reset the default compiler_so, which we might have changed for cuda 99 | self.compiler_so = default_compiler_so 100 | 101 | # inject our redefined _compile method into the class 102 | self._compile = _compile 103 | 104 | 105 | # run the customize_compiler 106 | class custom_build_ext(build_ext): 107 | def build_extensions(self): 108 | customize_compiler_for_nvcc(self.compiler) 109 | build_ext.build_extensions(self) 110 | 111 | ext_modules = [ 112 | Extension('poly_nms', 113 | ['poly_nms_kernel.cu', 'poly_nms.pyx'], 114 | library_dirs=[CUDA['lib64']], 115 | libraries=['cudart'], 116 | language='c++', 117 | runtime_library_dirs=[CUDA['lib64']], 118 | # this syntax is specific to this build system 119 | # we're only going to use certain compiler args with nvcc and not with 120 | # gcc the implementation of this trick is in customize_compiler() below 121 | extra_compile_args={'gcc': ["-Wno-unused-function"], 122 | 'nvcc': ['-arch=sm_35', 123 | '--ptxas-options=-v', 124 | '-c', 125 | '--compiler-options', 126 | "'-fPIC'"]}, 127 | include_dirs=[numpy_include, CUDA['include']] 128 | ), 129 | Extension('poly_overlaps', 130 | ['poly_overlaps_kernel.cu', 'poly_overlaps.pyx'], 131 | library_dirs=[CUDA['lib64']], 132 | libraries=['cudart'], 133 | language='c++', 134 | runtime_library_dirs=[CUDA['lib64']], 135 | # this syntax is specific to this build system 136 | # we're only going to use certain compiler args with nvcc and not with 137 | # gcc the implementation of this trick is in customize_compiler() below 138 | extra_compile_args={'gcc': ["-Wno-unused-function"], 139 | 'nvcc': ['-arch=sm_35', 140 | '--ptxas-options=-v', 141 | '-c', 142 | '--compiler-options', 143 | "'-fPIC'"]}, 144 | include_dirs=[numpy_include, CUDA['include']] 145 | ), 146 | ] 147 | setup( 148 | name='rotation', 149 | ext_modules=ext_modules, 150 | # inject our custom trigger 151 | cmdclass={'build_ext': custom_build_ext}, 152 | ) 153 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/polyiou.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | using namespace std; 8 | #define maxn 51 9 | const double eps=1E-8; 10 | int sig(double d){ 11 | return(d>eps)-(d<-eps); 12 | } 13 | struct Point{ 14 | double x,y; Point(){} 15 | Point(double x,double y):x(x),y(y){} 16 | bool operator==(const Point&p)const{ 17 | return sig(x-p.x)==0&&sig(y-p.y)==0; 18 | } 19 | }; 20 | double cross(Point o,Point a,Point b){ //叉积 21 | return(a.x-o.x)*(b.y-o.y)-(b.x-o.x)*(a.y-o.y); 22 | } 23 | double area(Point* ps,int n){ 24 | ps[n]=ps[0]; 25 | double res=0; 26 | for(int i=0;i0) pp[m++]=p[i]; 49 | // if(sig(cross(a,b,p[i]))!=sig(cross(a,b,p[i+1]))) 50 | // lineCross(a,b,p[i],p[i+1],pp[m++]); 51 | // } 52 | // n=0; 53 | // for(int i=0;i1&&p[n-1]==p[0])n--; 57 | //} 58 | void polygon_cut(Point*p,int&n,Point a,Point b, Point* pp){ 59 | // static Point pp[maxn]; 60 | int m=0;p[n]=p[0]; 61 | for(int i=0;i0) pp[m++]=p[i]; 63 | if(sig(cross(a,b,p[i]))!=sig(cross(a,b,p[i+1]))) 64 | lineCross(a,b,p[i],p[i+1],pp[m++]); 65 | } 66 | n=0; 67 | for(int i=0;i1&&p[n-1]==p[0])n--; 71 | } 72 | //---------------华丽的分隔线-----------------// 73 | //返回三角形oab和三角形ocd的有向交面积,o是原点// 74 | double intersectArea(Point a,Point b,Point c,Point d){ 75 | Point o(0,0); 76 | int s1=sig(cross(o,a,b)); 77 | int s2=sig(cross(o,c,d)); 78 | if(s1==0||s2==0)return 0.0;//退化，面积为0 79 | if(s1==-1) swap(a,b); 80 | if(s2==-1) swap(c,d); 81 | Point p[10]={o,a,b}; 82 | int n=3; 83 | Point pp[maxn]; 84 | polygon_cut(p,n,o,c, pp); 85 | polygon_cut(p,n,c,d, pp); 86 | polygon_cut(p,n,d,o, pp); 87 | double res=fabs(area(p,n)); 88 | if(s1*s2==-1) res=-res;return res; 89 | } 90 | //求两多边形的交面积 91 | double intersectArea(Point*ps1,int n1,Point*ps2,int n2){ 92 | if(area(ps1,n1)<0) reverse(ps1,ps1+n1); 93 | if(area(ps2,n2)<0) reverse(ps2,ps2+n2); 94 | ps1[n1]=ps1[0]; 95 | ps2[n2]=ps2[0]; 96 | double res=0; 97 | for(int i=0;i p, vector q) { 109 | Point ps1[maxn],ps2[maxn]; 110 | int n1 = 4; 111 | int n2 = 4; 112 | for (int i = 0; i < 4; i++) { 113 | ps1[i].x = p[i * 2]; 114 | ps1[i].y = p[i * 2 + 1]; 115 | 116 | ps2[i].x = q[i * 2]; 117 | ps2[i].y = q[i * 2 + 1]; 118 | } 119 | double inter_area = intersectArea(ps1, n1, ps2, n2); 120 | double union_area = fabs(area(ps1, n1)) + fabs(area(ps2, n2)) - inter_area; 121 | double iou = inter_area / union_area; 122 | 123 | // cout << "inter_area:" << inter_area << endl; 124 | // cout << "union_area:" << union_area << endl; 125 | // cout << "iou:" << iou << endl; 126 | 127 | return iou; 128 | } 129 | // 130 | //int main(){ 131 | // double p[8] = {0, 0, 1, 0, 1, 1, 0, 1}; 132 | // double q[8] = {0.5, 0.5, 1.5, 0.5, 1.5, 1.5, 0.5, 1.5}; 133 | // vector P(p, p + 8); 134 | // vector Q(q, q + 8); 135 | // iou_poly(P, Q); 136 | // return 0; 137 | //} 138 | 139 | //int main(){ 140 | // double p[8] = {0, 0, 1, 0, 1, 1, 0, 1}; 141 | // double q[8] = {0.5, 0.5, 1.5, 0.5, 1.5, 1.5, 0.5, 1.5}; 142 | // iou_poly(p, q); 143 | // return 0; 144 | //} -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/polyiou.h: -------------------------------------------------------------------------------- 1 | // 2 | // Created by dingjian on 18-2-3. 3 | // 4 | 5 | #ifndef POLYIOU_POLYIOU_H 6 | #define POLYIOU_POLYIOU_H 7 | 8 | #include 9 | double iou_poly(std::vector p, std::vector q); 10 | #endif //POLYIOU_POLYIOU_H 11 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/polyiou.i: -------------------------------------------------------------------------------- 1 | %module polyiou 2 | %include "std_vector.i" 3 | 4 | namespace std { 5 | %template(VectorDouble) vector; 6 | }; 7 | 8 | %{ 9 | #define SWIG_FILE_WITH_INIT 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | #include "polyiou.h" 16 | %} 17 | 18 | %include "polyiou.h" 19 | 20 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/prepare_dota1_ms.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | 4 | from DOTA_devkit.DOTA2JSON import generate_json_labels 5 | from DOTA_devkit.DOTA2COCO_poly import DOTA2COCOTrain, DOTA2COCOTest, wordname_15 6 | 7 | from DOTA_devkit.ImgSplit_multi_process import splitbase as splitbase_trainval 8 | from DOTA_devkit.SplitOnlyImage_multi_process import \ 9 | splitbase as splitbase_test 10 | 11 | 12 | def mkdir_if_not_exists(path): 13 | if not osp.exists(path): 14 | os.mkdir(path) 15 | 16 | def prepare_multi_scale_data(src_path, dst_path, gap=200, subsize=1024, scales=[0.5, 1.0, 1.5], num_process=32): 17 | """Prepare DOTA split data and labels 18 | Args: 19 | src_path: dataset path 20 | dst_path: output path 21 | gap: overlapping area 22 | subsize: size of chip image 23 | scales: multi-scale settings 24 | num_process: num of processer 25 | """ 26 | dst_train_path = osp.join(dst_path, 'train_split') 27 | dst_val_path = osp.join(dst_path, 'val_split') 28 | dst_trainval_path = osp.join(dst_path, 'trainval_split') 29 | dst_test_base_path = osp.join(dst_path, 'test_split') 30 | dst_test_path = osp.join(dst_path, 'test_split/images') 31 | # make dst path if not exist 32 | mkdir_if_not_exists(dst_path) 33 | mkdir_if_not_exists(dst_train_path) 34 | mkdir_if_not_exists(dst_val_path) 35 | mkdir_if_not_exists(dst_test_base_path) 36 | mkdir_if_not_exists(dst_test_path) 37 | # split train data 38 | print('split train data') 39 | split_train = splitbase_trainval(osp.join(src_path, 'train'), dst_train_path, 40 | gap=gap, subsize=subsize, num_process=num_process) 41 | for scale in scales: 42 | split_train.splitdata(scale) 43 | print('split val data') 44 | # split val data 45 | split_val = splitbase_trainval(osp.join(src_path, 'val'), dst_val_path, 46 | gap=gap, subsize=subsize, num_process=num_process) 47 | for scale in scales: 48 | split_val.splitdata(scale) 49 | # split test data 50 | print('split test data') 51 | split_test = splitbase_test(osp.join(src_path, 'test/images'), dst_test_path, 52 | gap=gap, subsize=subsize, num_process=num_process) 53 | for scale in scales: 54 | split_test.splitdata(scale) 55 | 56 | # prepare trainval data 57 | print('move train val to trainval') 58 | mkdir_if_not_exists(dst_trainval_path) 59 | os.system( 60 | 'mv {}/* {}'.format(dst_train_path, dst_trainval_path)) 61 | os.system('find '+dst_val_path+'/images/ -name "*.png" -exec mv {} ' + 62 | dst_trainval_path + '/images/ \\;') 63 | os.system('find '+dst_val_path+'/labelTxt/ -name "*.txt" -exec mv {} ' + 64 | dst_trainval_path + '/labelTxt/ \\;') 65 | 66 | print('generate labels with json format') 67 | generate_json_labels(dst_trainval_path, osp.join( 68 | dst_trainval_path, 'trainval.json')) 69 | generate_json_labels(dst_test_base_path, osp.join( 70 | dst_test_base_path, 'test.json'), trainval=False) 71 | print('generate labels with coco format') 72 | DOTA2COCOTrain(dst_trainval_path, 73 | osp.join(dst_trainval_path, 'trainval_coco_8point.json'), 74 | wordname_15) 75 | DOTA2COCOTest(dst_test_base_path, 76 | osp.join(dst_test_base_path, 'test_coco_8point.json'), 77 | wordname_15) 78 | 79 | 80 | if __name__ == '__main__': 81 | # single scale 82 | prepare_multi_scale_data('/data1/dataset_demo/DOTA_demo/', 83 | '/data1/OrientedRepPoints/data/dota_1024', scales=[1.0], gap=200) 84 | # multi scale 85 | # prepare_multi_scale_data('/mnt/SSD/lwt_workdir/data/dota_new/', 86 | # '/mnt/SSD/lwt_workdir/data/dota_1024_ms', scales=[0.5, 1.0, 1.5], gap=500) 87 | print('done') 88 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/prepare_hrsc2016.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | from DOTA_devkit.HRSC2DOTA import generate_txt_labels 4 | from DOTA_devkit.DOTA2JSON import generate_json_labels 5 | 6 | def preprare_hrsc2016(data_dir): 7 | train_dir = osp.join(data_dir,'Train') 8 | test_dir = osp.join(data_dir, 'Test') 9 | # convert hrsc2016 to dota raw format 10 | generate_txt_labels(train_dir) 11 | generate_txt_labels(test_dir) 12 | # convert it to json format 13 | generate_json_labels(train_dir,osp.join(train_dir,'trainval.json')) 14 | generate_json_labels(test_dir,osp.join(test_dir,'test.json'), trainval=False) 15 | 16 | if __name__ == '__main__': 17 | hrsc2016_dir = '/mnt/SSD/lwt_workdir/BeyondBoundingBox/data/hrsc2016/' 18 | preprare_hrsc2016(hrsc2016_dir) 19 | print('done') 20 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/results_ensemble.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import shutil 4 | 5 | def GetFileFromThisRootDir(dir,ext = None): 6 | allfiles = [] 7 | needExtFilter = (ext != None) 8 | for root,dirs,files in os.walk(dir): 9 | for filespath in files: 10 | filepath = os.path.join(root, filespath) 11 | extension = os.path.splitext(filepath)[1][1:] 12 | if needExtFilter and extension in ext: 13 | allfiles.append(filepath) 14 | elif not needExtFilter: 15 | allfiles.append(filepath) 16 | return allfiles 17 | 18 | def custombasename(fullname): 19 | return os.path.basename(os.path.splitext(fullname)[0]) 20 | 21 | def results_ensemble(srcpath_1, srcpath_2, dstpath): 22 | """ 23 | 将srcpath_1,srcpath_2文件夹中的所有txt中的目标提取出来, 并叠加在一起存入 dstpath 24 | """ 25 | if os.path.exists(dstpath): 26 | shutil.rmtree(dstpath) # delete output folderX 27 | os.makedirs(dstpath) 28 | 29 | filelist_1 = GetFileFromThisRootDir(srcpath_1) # srcpath文件夹下的所有文件相对路径 eg:['Task1_??.txt', ..., '?.txt'] 30 | filelist_2 = GetFileFromThisRootDir(srcpath_2) 31 | for index, fullname_1 in enumerate(filelist_1): # Task1_??.txt' 32 | fullname_2 = filelist_2[index] 33 | basename = custombasename(fullname_1) # 只留下文件名 eg:'Task1_??' 34 | dstname = os.path.join(dstpath, basename + '.txt') # eg: ./.../Task1_plane.txt 35 | 36 | with open(dstname, 'a') as f_out: 37 | # merge first txt 38 | with open(fullname_1, 'r') as f1: 39 | lines = f1.readlines() 40 | for line in lines: 41 | f_out.writelines(line) 42 | # merge second txt 43 | with open(fullname_2, 'r') as f2: 44 | lines = f2.readlines() 45 | for line in lines: 46 | f_out.writelines(line) 47 | pass 48 | 49 | def parse_args(): 50 | parser = argparse.ArgumentParser(description='model ensemble') 51 | parser.add_argument('--srcpath_1', default='/OrientedRepPoints/tools/parse_pkl/evaluation_results/ORR_results/', help='srcpath_1') 52 | parser.add_argument('--srcpath_2', default='/OrientedRepPoints/tools/parse_pkl/evaluation_results/ROI_results/', help='srcpath_2') 53 | parser.add_argument('--dstpath', default='/OrientedRepPoints/tools/parse_pkl/evaluation_results/orientedreppoints_ROIRT_ensemble/', help='dstpath') 54 | args = parser.parse_args() 55 | return args 56 | 57 | def main(): 58 | args = parse_args() 59 | srcpath_1 = args.srcpath_1 60 | srcpath_2 = args.srcpath_2 61 | dstpath = args.dstpath 62 | 63 | results_ensemble(srcpath_1, srcpath_2, dstpath) 64 | 65 | 66 | if __name__ == '__main__': 67 | main() -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/results_obb2hbb.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import shutil 4 | 5 | def parse_args(): 6 | parser = argparse.ArgumentParser(description='Train a detector') 7 | parser.add_argument('--srcpath', default=r'/OrientedRepPoints/tools/parse_pkl/evaluation_results/OBB_results/') 8 | parser.add_argument('--dstpath', default='/OrientedRepPoints/tools/parse_pkl/evaluation_results/HBB_results/', 9 | help='dota version') 10 | args = parser.parse_args() 11 | 12 | return args 13 | 14 | 15 | def GetFileFromThisRootDir(dir,ext = None): 16 | allfiles = [] 17 | needExtFilter = (ext != None) 18 | for root,dirs,files in os.walk(dir): 19 | for filespath in files: 20 | filepath = os.path.join(root, filespath) 21 | extension = os.path.splitext(filepath)[1][1:] 22 | if needExtFilter and extension in ext: 23 | allfiles.append(filepath) 24 | elif not needExtFilter: 25 | allfiles.append(filepath) 26 | return allfiles 27 | 28 | def custombasename(fullname): 29 | return os.path.basename(os.path.splitext(fullname)[0]) 30 | 31 | def OBB2HBB(srcpath, dstpath): 32 | filenames = GetFileFromThisRootDir(srcpath) 33 | if os.path.exists(dstpath): 34 | shutil.rmtree(dstpath) # delete output folderX 35 | os.makedirs(dstpath) 36 | 37 | for file in filenames: # eg: /.../task1_plane.txt 38 | 39 | basename = custombasename(file) # 只留下文件名 eg:'task1_plane' 40 | class_basename = basename.split('_')[-1] 41 | with open(file, 'r') as f_in: 42 | with open(os.path.join(dstpath, 'Task2_' + class_basename + '.txt'), 'w') as f_out: 43 | lines = f_in.readlines() 44 | splitlines = [x.strip().split() for x in lines] # list: n*[] 45 | for index, splitline in enumerate(splitlines): 46 | imgname = splitline[0] 47 | score = splitline[1] 48 | poly = splitline[2:] 49 | poly = list(map(float, poly)) 50 | xmin, xmax, ymin, ymax = min(poly[0::2]), max(poly[0::2]), min(poly[1::2]), max(poly[1::2]) 51 | rec_poly = [xmin, ymin, xmax, ymax] 52 | outline = imgname + ' ' + score + ' ' + ' '.join(map(str, rec_poly)) 53 | if index != (len(splitlines) - 1): 54 | outline = outline + '\n' 55 | f_out.write(outline) 56 | 57 | if __name__ == '__main__': 58 | args = parse_args() 59 | OBB2HBB(args.srcpath, args.dstpath) 60 | -------------------------------------------------------------------------------- /yolov5_obb/DOTA_devkit/setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | setup.py file for SWIG example 3 | """ 4 | from distutils.core import setup, Extension 5 | import numpy 6 | 7 | polyiou_module = Extension('_polyiou', 8 | sources=['polyiou_wrap.cxx', 'polyiou.cpp'], 9 | ) 10 | setup(name = 'polyiou', 11 | version = '0.1', 12 | author = "SWIG Docs", 13 | description = """Simple swig example from docs""", 14 | ext_modules = [polyiou_module], 15 | py_modules = ["polyiou"], 16 | ) 17 | -------------------------------------------------------------------------------- /yolov5_obb/README.md: -------------------------------------------------------------------------------- 1 | # Yolov5 for Oriented Object Detection 2 | ![图片](./docs/detection.png) 3 | ![train_batch0.jpg](./docs/train_batch6.jpg) 4 | ![results.png](./docs/results.png) 5 | 6 | The code for the implementation of “[Yolov5](https://github.com/ultralytics/yolov5) + [Circular Smooth Label](https://arxiv.org/abs/2003.05597v2)”. 7 | 8 | # Results and Models 9 | The results on **DOTA_subsize1024_gap200_rate1.0** test-dev set are shown in the table below. (**password: yolo**) 10 | 11 | |Model
(link) |Size
(pixels) | TTA
(multi-scale/
rotate testing) | OBB mAPtest
0.5
DOTAv1.0 | OBB mAPtest
0.5
DOTAv1.5 | OBB mAPtest
0.5
DOTAv2.0 | Speed
CPU b1
(ms)|Speed
2080Ti b1
(ms) |Speed
2080Ti b16
(ms) |params
(M) |FLOPs
@1024 (B) 12 | | ---- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- 13 | |yolov5m [[baidu](https://pan.baidu.com/s/1UPNaMuQ_gNce9167FZx8-w)/[google](https://drive.google.com/file/d/1NMgxcN98cmBg9_nVK4axxqfiq4pYh-as/view?usp=sharing)] |1024 | × |**77.30** |**73.19** |**58.01** |**328.2** |**16.9** |**11.3** |**21.6** |**50.5** 14 | |[yolov5s]() |1024 | × |**76.79** |- |- |- |**15.6** | - |**7.54** |**17.5** 15 | |[yolov5tiny]() |1024 | × |**73.26** |- |- |- |**15.2** | - |**2.02** |**5.0** 16 | 17 | 18 |
19 | Table Notes (click to expand / **点我看更多**) 20 | 21 | * All checkpoints are trained to 300 epochs with [COCO pre-trained checkpoints](https://github.com/ultralytics/yolov5/releases/tag/v6.0), default settings and hyperparameters. 22 | * **mAPtest dota** values are for single-model single-scale on [DOTA](https://captain-whu.github.io/DOTA/index.html)(1024,1024,200,1.0) dataset.
Reproduce Example: 23 | ```shell 24 | python val.py --data 'data/dotav15_poly.yaml' --img 1024 --conf 0.01 --iou 0.4 --task 'test' --batch 16 --save-json --name 'dotav15_test_split' 25 | python tools/TestJson2VocClassTxt.py --json_path 'runs/val/dotav15_test_split/best_obb_predictions.json' --save_path 'runs/val/dotav15_test_split/obb_predictions_Txt' 26 | python DOTA_devkit/ResultMerge_multi_process.py --scrpath 'runs/val/dotav15_test_split/obb_predictions_Txt' --dstpath 'runs/val/dotav15_test_split/obb_predictions_Txt_Merged' 27 | zip the poly format results files and submit it to https://captain-whu.github.io/DOTA/evaluation.html 28 | ``` 29 | * **Speed** averaged over DOTAv1.5 val_split_subsize1024_gap200 images using a 2080Ti gpu. NMS + pre-process times is included.
Reproduce by `python val.py --data 'data/dotav15_poly.yaml' --img 1024 --task speed --batch 1` 30 | 31 | 32 |
33 | 34 | # [Updates](./docs/ChangeLog.md) 35 | - [2022/1/7] : **Faster and stronger**, some bugs fixed, yolov5 base version updated. 36 | 37 | 38 | # Installation 39 | Please refer to [install.md](./docs/install.md) for installation and dataset preparation. 40 | 41 | 42 | # Getting Started 43 | This repo is based on [yolov5](https://github.com/ultralytics/yolov5). 44 | 45 | And this repo has been rebuilt, Please see [GetStart.md](./docs/GetStart.md) for the Oriented Detection latest basic usage. 46 | 47 | 48 | ## Acknowledgements 49 | I have used utility functions from other wonderful open-source projects. Espeicially thank the authors of: 50 | 51 | * [ultralytics/yolov5](https://github.com/ultralytics/yolov5). 52 | * [Thinklab-SJTU/CSL_RetinaNet_Tensorflow](https://github.com/Thinklab-SJTU/CSL_RetinaNet_Tensorflow). 53 | * [jbwang1997/OBBDetection](https://github.com/jbwang1997/OBBDetection) 54 | * [CAPTAIN-WHU/DOTA_devkit](https://github.com/CAPTAIN-WHU/DOTA_devkit) 55 | 56 | 57 | ## More detailed explanation 58 | 想要了解相关实现的细节和原理可以看我的知乎文章: 59 | 60 | * [自己改建YOLOv5旋转目标的踩坑记录](https://www.zhihu.com/column/c_1358464959123390464). 61 | 62 | ## 有问题反馈 63 | 在使用中有任何问题，建议先按照install.md检查环境依赖项，再按照getstart.md检查使用流程是否正确，善用搜索引擎和github中的issue搜索框，可以极大程度上节省你的时间。 64 | 若遇到的是新问题，可以用以下联系方式跟我交流 65 | 66 | * 知乎（@[略略略](https://www.zhihu.com/people/lue-lue-lue-3-92-86)） 67 | * 代码问题提issues,其他问题请知乎上联系 68 | 69 | ## 关于作者 70 | 71 | ```javascript 72 | Name : "胡凯旋" 73 | describe myself："咸鱼一枚" 74 | 75 | -------------------------------------------------------------------------------- /yolov5_obb/data/dotav15_poly.yaml: -------------------------------------------------------------------------------- 1 | 2 | 3 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 4 | #path: /media/test/DOTA/DOTAv1.5 # dataset root dir 5 | #train: train_split_1024_gap200/images # train images (relative to 'path') 6 | #val: val_split_1024_gap200/images # val images (relative to 'path') 7 | #test: test_split_1024/images # test images (optional) 8 | 9 | # Classes 10 | #nc: 2 # number of classes 11 | #names: ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 12 | # 'large-vehicle', 'ship', 'tennis-court', 'basketball-court', 'storage-tank', 13 | # 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter', 14 | # 'container-crane'] # class names 15 | path: data/split/ # dataset root dir 16 | train: images/ # train images (relative to 'path') 17 | val: images/ # val images (relative to 'path') 18 | test: images/ # test images (optional) 19 | nc: 1 # number of classes 20 | names: ['tomato'] 21 | 22 | -------------------------------------------------------------------------------- /yolov5_obb/data/hyps/obb/hyp.finetune_dota.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 4 | lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf) 5 | momentum: 0.937 # SGD momentum/Adam beta1 6 | weight_decay: 0.0005 # optimizer weight decay 5e-4 7 | warmup_epochs: 3.0 # warmup epochs (fractions ok) 8 | warmup_momentum: 0.8 # warmup initial momentum 9 | warmup_bias_lr: 0.1 # warmup initial bias lr 10 | box: 0.05 # box loss gain 11 | cls: 0.5 # cls loss gain 12 | cls_pw: 1.0 # cls BCELoss positive_weight 13 | theta: 0.5 14 | theta_pw: 1.0 15 | obj: 1.0 # obj loss gain (scale with pixels) 16 | obj_pw: 1.0 # obj BCELoss positive_weight 17 | iou_t: 0.2 # IoU training threshold 18 | anchor_t: 4.0 # anchor-multiple threshold 19 | # anchors: 3.2 # anchors per output layer (0 to ignore) 20 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 21 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 22 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 23 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 24 | degrees: 180.0 # image rotation (+/- deg) 25 | translate: 0.1 # image translation (+/- fraction) 26 | scale: 0.5 # image scale (+/- gain) 27 | shear: 0.0 # image shear (+/- deg) 28 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 29 | flipud: 0.5 # image flip up-down (probability) 30 | fliplr: 0.5 # image flip left-right (probability) 31 | mosaic: 0.95 # image mosaic (probability) 32 | mixup: 0.1 # image mixup (probability) 33 | copy_paste: 0.0 # segment copy-paste (probability) 34 | cls_theta: 180 35 | csl_radius: 2.0 36 | -------------------------------------------------------------------------------- /yolov5_obb/data/hyps/obb/hyp.finetune_dota_CloseAug.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.00258 4 | lrf: 0.17 5 | momentum: 0.779 6 | weight_decay: 0.00058 7 | warmup_epochs: 1.33 8 | warmup_momentum: 0.86 9 | warmup_bias_lr: 0.0711 10 | box: 0.0539 11 | cls: 0.299 12 | cls_pw: 0.825 13 | theta: 0.299 14 | theta_pw: 0.825 15 | obj: 0.632 16 | obj_pw: 1.0 17 | iou_t: 0.2 18 | anchor_t: 3.44 19 | anchors: 3.2 20 | fl_gamma: 0.0 21 | hsv_h: 0.0188 22 | hsv_s: 0.704 23 | hsv_v: 0.36 24 | degrees: 180.0 25 | translate: 0.0902 26 | scale: 0.25 27 | shear: 0.0 28 | perspective: 0.0 29 | flipud: 0.5 30 | fliplr: 0.5 31 | mosaic: 0.0 32 | mixup: 0.0 33 | copy_paste: 0.0 34 | cls_theta: 180 35 | csl_radius: 2.0 -------------------------------------------------------------------------------- /yolov5_obb/data/scripts/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 3 | # Download latest models from https://github.com/ultralytics/yolov5/releases 4 | # Example usage: bash path/to/download_weights.sh 5 | # parent 6 | # └── yolov5 7 | # ├── yolov5s.pt ← downloads here 8 | # ├── yolov5m.pt 9 | # └── ... 10 | 11 | python - <test
@0.5| fps | 9 | ---- | ----- | ------ | ----- | ----- | ----- | ----- | ----- | 10 | yolov5m-old | DOTAv1.5train_subsize1024_gap200_rate1.0|75 |300 |2080Ti |96h |68.36 |20 | 11 | **yolov5m-new**| DOTAv1.5train_subsize1024_gap200_rate1.0|75 |300 |2080Ti |**15h**|**73.19**|**59** | 12 | 13 | 3. Some Bugs Fixed. 14 | 15 | |Bug | Fixed | Describe 16 | |---- |------ | ------ 17 | |Don't support validation | ✔ | Support hbb validation in training, which is faster than obb validation| 18 | |Don't support single class training | ✔ | But it will get weaker results than **'nc=2'** | 19 | |Image must meets Height=Width | ✔ | - | 20 | 21 | 4. support obb_nms gpu version. 22 | -------------------------------------------------------------------------------- /yolov5_obb/docs/GetStart.md: -------------------------------------------------------------------------------- 1 | # Getting Started 2 | 3 | This page provides basic usage about yolov5-obb. For installation instructions, please see [install.md](./install.md). 4 | 5 | # Train a model 6 | 7 | **1. Prepare custom dataset files** 8 | 9 | 1.1 Make sure the labels format is [poly classname diffcult], e.g., 10 | ``` 11 | x1 y1 x2 y2 x3 y3 x4 y4 classname diffcult 12 | 13 | 1686.0 1517.0 1695.0 1511.0 1711.0 1535.0 1700.0 1541.0 large-vehicle 1 14 | ``` 15 | **(*Note: You can set **diffcult=0**)** 16 | ![image](https://user-images.githubusercontent.com/72599120/159213229-b7c2fc5c-b140-4f10-9af8-2cbc405b0cd3.png) 17 | 18 | 19 | 1.2 Split the dataset. 20 | ```shell 21 | cd yolov5_obb 22 | python DOTA_devkit/ImgSplit_multi_process.py 23 | ``` 24 | or Use the orignal dataset. 25 | ```shell 26 | cd yolov5_obb 27 | ``` 28 | **(*Note: High resolution image dataset needs to be splited to get better performance in small objects)** 29 | 30 | 31 | **2. Train** 32 | 33 | 2.1 Train with specified GPUs. (for example with GPU=3) 34 | 35 | ```shell 36 | python train.py --device 3 37 | ``` 38 | 39 | 2.2 Train with multiple(4) GPUs. (DDP Mode) 40 | 41 | ```shell 42 | python -m torch.distributed.launch --nproc_per_node 4 train.py --device 0,1,2,3 43 | ``` 44 | 45 | 2.3 Train the orignal dataset demo. 46 | ```shell 47 | python train.py --data 'data/yolov5obb_demo.yaml' --epochs 10 --batch-size 1 --img 1024 --device 0 48 | ``` 49 | 50 | 2.4 Train the splited dataset demo. 51 | ```shell 52 | python train.py --data 'data/yolov5obb_demo_split.yaml' --epochs 10 --batch-size 2 --img 1024 --device 0 53 | ``` 54 | 55 | # Inferenece with pretrained models. (Splited Dataset) 56 | This repo provides the validation/testing scripts to evaluate the trained model. 57 | 58 | Examples: 59 | 60 | Assume that you have already downloaded the checkpoints to `runs/train/yolov5m_finetune/weights`. 61 | 62 | 1. Test yolov5-obb with single GPU. Get the HBB metrics. 63 | 64 | ```shell 65 | python val.py --task 'val' --device 0 --save-json --batch-size 2 --data 'data/yolov5obb_demo_split.yaml' --name 'obb_demo_split' 66 | 67 | Class Images Labels P R mAP@.5 mAP@.5:.95: 100%|██████████| 3/3 [00:02<00:00, 1.09it/s] 68 | all 6 68 0.921 0.914 0.966 0.776 69 | plane 6 16 0.946 1 0.995 0.934 70 | small-vehicle 6 35 0.928 0.741 0.916 0.599 71 | large-vehicle 6 17 0.89 1 0.986 0.793 72 | Speed: .................................................... per image at shape (2, 3, 1024, 1024) 73 | ... 74 | Evaluating pycocotools mAP... saving runs/val/obb_demo_split/best_obb_predictions.json... 75 | ---------------------The hbb and obb results has been saved in json file----------------------- 76 | ``` 77 | 78 | 2. Parse the results. Get the poly format results. 79 | ```shell 80 | python tools/TestJson2VocClassTxt.py --json_path 'runs/val/obb_demo_split/best_obb_predictions.json' --save_path 'runs/val/obb_demo_split/obb_predictions_Txt' 81 | ``` 82 | 83 | 3. Merge the results. (If you split your dataset) 84 | ```shell 85 | python DOTA_devkit/ResultMerge_multi_process.py \ 86 | --scrpath 'runs/val/obb_demo_split/obb_predictions_Txt' \ 87 | --dstpath 'runs/val/obb_demo_split/obb_predictions_Txt_Merged' 88 | ``` 89 | 90 | 4. Get the OBB metrics 91 | ```shell 92 | python DOTA_devkit/dota_evaluation_task1.py \ 93 | --detpath 'runs/val/obb_demo_split/obb_predictions_Txt_Merged/Task1_{:s}.txt' \ 94 | --annopath 'dataset/dataset_demo/labelTxt/{:s}.txt' \ 95 | --imagesetfile 'dataset/dataset_demo/imgnamefile.txt' 96 | 97 | ... 98 | map: 0.6666666666666669 99 | classaps: [100. 0. 100.] 100 | ``` 101 | 102 | # Inferenece with pretrained models. (Original Dataset) 103 | We provide the validation/testing scripts to evaluate the trained model. 104 | 105 | Examples: 106 | 107 | Assume that you have already downloaded the checkpoints to `runs/train/yolov5m_finetune/weights`. 108 | 109 | 1. Test yolov5-obb with single GPU. Get the HBB metrics. 110 | 111 | ```shell 112 | python val.py --task 'val' --device 0 --save-json --batch-size 1 --data 'data/yolov5obb_demo.yaml' --name 'obb_demo' --img 2048 113 | 114 | Class Images Labels P R mAP@.5 mAP@.5:.95: 100%|██████████| 1/1 [00:00<00:00, 1.98it/s] 115 | all 1 56 0.97 0.85 0.953 0.752 116 | plane 1 11 1 1 0.995 0.944 117 | small-vehicle 1 34 1 0.641 0.889 0.535 118 | large-vehicle 1 11 0.91 0.909 0.976 0.777 119 | Speed: .................................................... per image at shape (1, 3, 2048, 2048) 120 | ... 121 | Evaluating pycocotools mAP... saving runs/val/obb_demo/best_obb_predictions.json... 122 | ---------------------The hbb and obb results has been saved in json file----------------------- 123 | ``` 124 | 125 | 2. Parse the results. Get the poly format results. 126 | ```shell 127 | python tools/TestJson2VocClassTxt.py --json_path 'runs/val/obb_demo/best_obb_predictions.json' --save_path 'runs/val/obb_demo/obb_predictions_Txt' 128 | ``` 129 | 130 | 3. Get the OBB metrics 131 | ```shell 132 | python DOTA_devkit/dota_evaluation_task1.py \ 133 | --detpath 'runs/val/obb_demo/obb_predictions_Txt/Task1_{:s}.txt' \ 134 | --annopath 'dataset/dataset_demo/labelTxt/{:s}.txt' \ 135 | --imagesetfile 'dataset/dataset_demo/imgnamefile.txt' 136 | 137 | ... 138 | map: 0.6666666666666669 139 | classaps: [100. 0. 100.] 140 | ``` 141 | 142 | # Run inference on images, videos, directories, streams, etc. Then save the detection file. 143 | 1. image demo 144 | ```shell 145 | python detect.py --weights 'runs/train/yolov5m_finetune/weights/best.pt' \ 146 | --source 'dataset/dataset_demo/images/' \ 147 | --img 2048 --device 2 --hide-labels --hide-conf 148 | ``` 149 | 150 | ***If you want to evaluate the result on DOTA test-dev, please zip the poly format results files and submit it to the [evaluation server](https://captain-whu.github.io/DOTA/index.html).** 151 | -------------------------------------------------------------------------------- /yolov5_obb/docs/detection.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/docs/detection.png -------------------------------------------------------------------------------- /yolov5_obb/docs/install.md: -------------------------------------------------------------------------------- 1 | # INSTAllation 2 | ## Requirements 3 | * Windows **(not Recommend)**, Linux **(Recommend)** 4 | * Python 3.7+ 5 | * PyTorch ≥ 1.7 6 | * CUDA 9.0 or higher 7 | 8 | I have tested the following versions of OS and softwares： 9 | * OS：Ubuntu 16.04/18.04 10 | * CUDA: 10.0/10.1/10.2/11.3 11 | 12 | ## Install 13 | **CUDA Driver Version ≥ CUDA Toolkit Version(runtime version) = torch.version.cuda** 14 | 15 | a. Create a conda virtual environment and activate it, e.g., 16 | ``` 17 | conda create -n Py39_Torch1.10_cu11.3 python=3.9 -y 18 | source activate Py39_Torch1.10_cu11.3 19 | ``` 20 | b. Make sure your CUDA runtime api version ≤ CUDA driver version. (for example 11.3 ≤ 11.4) 21 | ``` 22 | nvcc -V 23 | nvidia-smi 24 | ``` 25 | c. Install PyTorch and torchvision following the [official instructions](https://pytorch.org/), Make sure cudatoolkit version same as CUDA runtime api version, e.g., 26 | ``` 27 | pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html 28 | nvcc -V 29 | python 30 | >>> import torch 31 | >>> torch.version.cuda 32 | >>> exit() 33 | ``` 34 | d. Clone the yolov5-obb repository. 35 | ``` 36 | git clone https://github.com/hukaixuan19970627/yolov5_obb.git 37 | cd yolov5_obb 38 | ``` 39 | e. Install yolov5-obb. 40 | 41 | ```python 42 | pip install -r requirements.txt 43 | cd utils/nms_rotated 44 | python setup.py develop #or "pip install -v -e ." 45 | ``` 46 | 47 | ## Install DOTA_devkit. 48 | **(Custom Install, it's just a tool to split the high resolution image and evaluation the obb)** 49 | 50 | ``` 51 | cd yolov5_obb/DOTA_devkit 52 | sudo apt-get install swig 53 | swig -c++ -python polyiou.i 54 | python setup.py build_ext --inplace 55 | ``` 56 | 57 | ## Prepare dataset 58 | ``` 59 | parent 60 | ├── yolov5 61 | └── datasets 62 | └── DOTAv1.5 63 | ├── train_split_rate1.0_subsize1024_gap200 64 | ├── train_split_rate1.0_subsize1024_gap200 65 | └── test_split_rate1.0_subsize1024_gap200 66 | ├── images 67 | └── labelTxt 68 | 69 | ``` 70 | 71 | **Note:** 72 | * DOTA is a high resolution image dataset, so it must be splited before training to get a better performance. 73 | -------------------------------------------------------------------------------- /yolov5_obb/docs/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/docs/results.png -------------------------------------------------------------------------------- /yolov5_obb/docs/train_batch6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/docs/train_batch6.jpg -------------------------------------------------------------------------------- /yolov5_obb/gen_wts.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import argparse 3 | import os 4 | import struct 5 | import torch 6 | from utils.torch_utils import select_device 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser(description='Convert .pt file to .wts') 11 | parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)') 12 | parser.add_argument('-o', '--output', help='Output (.wts) file path (optional)') 13 | args = parser.parse_args() 14 | if not os.path.isfile(args.weights): 15 | raise SystemExit('Invalid input file') 16 | if not args.output: 17 | args.output = os.path.splitext(args.weights)[0] + '.wts' 18 | elif os.path.isdir(args.output): 19 | args.output = os.path.join( 20 | args.output, 21 | os.path.splitext(os.path.basename(args.weights))[0] + '.wts') 22 | return args.weights, args.output 23 | 24 | 25 | pt_file, wts_file = parse_args() 26 | 27 | # Initialize 28 | device = select_device('cpu') 29 | # Load model 30 | model = torch.load(pt_file, map_location=device) # load to FP32 31 | model = model['ema' if model.get('ema') else 'model'].float() 32 | 33 | # update anchor_grid info 34 | anchor_grid = model.model[-1].anchors * model.model[-1].stride[...,None,None] 35 | # model.model[-1].anchor_grid = anchor_grid 36 | delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer 37 | model.model[-1].register_buffer("anchor_grid",anchor_grid) #The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight. 38 | 39 | model.to(device).eval() 40 | 41 | with open(wts_file, 'w') as f: 42 | f.write('{}\n'.format(len(model.state_dict().keys()))) 43 | for k, v in model.state_dict().items(): 44 | vr = v.reshape(-1).cpu().numpy() 45 | f.write('{} {} '.format(k, len(vr))) 46 | for vv in vr: 47 | f.write(' ') 48 | f.write(struct.pack('>f' ,float(vv)).hex()) 49 | f.write('\n') 50 | -------------------------------------------------------------------------------- /yolov5_obb/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__init__.py -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/common.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/common.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/common.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/common.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/experimental.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/experimental.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/experimental.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/experimental.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/yolo.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/yolo.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/__pycache__/yolo.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/models/__pycache__/yolo.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/models/experimental.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Experimental modules 4 | """ 5 | import math 6 | 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | 11 | from models.common import Conv 12 | from utils.downloads import attempt_download 13 | 14 | 15 | class CrossConv(nn.Module): 16 | # Cross Convolution Downsample 17 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 18 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 19 | super().__init__() 20 | c_ = int(c2 * e) # hidden channels 21 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 22 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 23 | self.add = shortcut and c1 == c2 24 | 25 | def forward(self, x): 26 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 27 | 28 | 29 | class Sum(nn.Module): 30 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 31 | def __init__(self, n, weight=False): # n: number of inputs 32 | super().__init__() 33 | self.weight = weight # apply weights boolean 34 | self.iter = range(n - 1) # iter object 35 | if weight: 36 | self.w = nn.Parameter(-torch.arange(1.0, n) / 2, requires_grad=True) # layer weights 37 | 38 | def forward(self, x): 39 | y = x[0] # no weight 40 | if self.weight: 41 | w = torch.sigmoid(self.w) * 2 42 | for i in self.iter: 43 | y = y + x[i + 1] * w[i] 44 | else: 45 | for i in self.iter: 46 | y = y + x[i + 1] 47 | return y 48 | 49 | 50 | class MixConv2d(nn.Module): 51 | # Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595 52 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kernel, stride, ch_strategy 53 | super().__init__() 54 | n = len(k) # number of convolutions 55 | if equal_ch: # equal c_ per group 56 | i = torch.linspace(0, n - 1E-6, c2).floor() # c2 indices 57 | c_ = [(i == g).sum() for g in range(n)] # intermediate channels 58 | else: # equal weight.numel() per group 59 | b = [c2] + [0] * n 60 | a = np.eye(n + 1, n, k=-1) 61 | a -= np.roll(a, 1, axis=1) 62 | a *= np.array(k) ** 2 63 | a[0] = 1 64 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 65 | 66 | self.m = nn.ModuleList( 67 | [nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)]) 68 | self.bn = nn.BatchNorm2d(c2) 69 | self.act = nn.SiLU() 70 | 71 | def forward(self, x): 72 | return self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 73 | 74 | 75 | class Ensemble(nn.ModuleList): 76 | # Ensemble of models 77 | def __init__(self): 78 | super().__init__() 79 | 80 | def forward(self, x, augment=False, profile=False, visualize=False): 81 | y = [] 82 | for module in self: 83 | y.append(module(x, augment, profile, visualize)[0]) 84 | # y = torch.stack(y).max(0)[0] # max ensemble 85 | # y = torch.stack(y).mean(0) # mean ensemble 86 | y = torch.cat(y, 1) # nms ensemble 87 | return y, None # inference, train output 88 | 89 | 90 | def attempt_load(weights, map_location=None, inplace=True, fuse=True): 91 | from models.yolo import Detect, Model 92 | 93 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 94 | model = Ensemble() 95 | for w in weights if isinstance(weights, list) else [weights]: 96 | ckpt = torch.load(attempt_download(w), map_location=map_location) # load 97 | if fuse: 98 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model 99 | else: 100 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval()) # without layer fuse 101 | 102 | # Compatibility updates 103 | for m in model.modules(): 104 | if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model]: 105 | m.inplace = inplace # pytorch 1.7.0 compatibility 106 | if type(m) is Detect: 107 | if not isinstance(m.anchor_grid, list): # new Detect Layer compatibility 108 | delattr(m, 'anchor_grid') 109 | setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl) 110 | elif type(m) is Conv: 111 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 112 | 113 | if len(model) == 1: 114 | return model[-1] # return model 115 | else: 116 | print(f'Ensemble created with {weights}\n') 117 | for k in ['names']: 118 | setattr(model, k, getattr(model[-1], k)) 119 | model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride 120 | return model # return ensemble 121 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/anchors.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Default anchors for COCO data 3 | 4 | 5 | # P5 ------------------------------------------------------------------------------------------------------------------- 6 | # P5-640: 7 | anchors_p5_640: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | 13 | # P6 ------------------------------------------------------------------------------------------------------------------- 14 | # P6-640: thr=0.25: 0.9964 BPR, 5.54 anchors past thr, n=12, img_size=640, metric_all=0.281/0.716-mean/best, past_thr=0.469-mean: 9,11, 21,19, 17,41, 43,32, 39,70, 86,64, 65,131, 134,130, 120,265, 282,180, 247,354, 512,387 15 | anchors_p6_640: 16 | - [9,11, 21,19, 17,41] # P3/8 17 | - [43,32, 39,70, 86,64] # P4/16 18 | - [65,131, 134,130, 120,265] # P5/32 19 | - [282,180, 247,354, 512,387] # P6/64 20 | 21 | # P6-1280: thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1280, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 19,27, 44,40, 38,94, 96,68, 86,152, 180,137, 140,301, 303,264, 238,542, 436,615, 739,380, 925,792 22 | anchors_p6_1280: 23 | - [19,27, 44,40, 38,94] # P3/8 24 | - [96,68, 86,152, 180,137] # P4/16 25 | - [140,301, 303,264, 238,542] # P5/32 26 | - [436,615, 739,380, 925,792] # P6/64 27 | 28 | # P6-1920: thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1920, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 28,41, 67,59, 57,141, 144,103, 129,227, 270,205, 209,452, 455,396, 358,812, 653,922, 1109,570, 1387,1187 29 | anchors_p6_1920: 30 | - [28,41, 67,59, 57,141] # P3/8 31 | - [144,103, 129,227, 270,205] # P4/16 32 | - [209,452, 455,396, 358,812] # P5/32 33 | - [653,922, 1109,570, 1387,1187] # P6/64 34 | 35 | 36 | # P7 ------------------------------------------------------------------------------------------------------------------- 37 | # P7-640: thr=0.25: 0.9962 BPR, 6.76 anchors past thr, n=15, img_size=640, metric_all=0.275/0.733-mean/best, past_thr=0.466-mean: 11,11, 13,30, 29,20, 30,46, 61,38, 39,92, 78,80, 146,66, 79,163, 149,150, 321,143, 157,303, 257,402, 359,290, 524,372 38 | anchors_p7_640: 39 | - [11,11, 13,30, 29,20] # P3/8 40 | - [30,46, 61,38, 39,92] # P4/16 41 | - [78,80, 146,66, 79,163] # P5/32 42 | - [149,150, 321,143, 157,303] # P6/64 43 | - [257,402, 359,290, 524,372] # P7/128 44 | 45 | # P7-1280: thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1280, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 19,22, 54,36, 32,77, 70,83, 138,71, 75,173, 165,159, 148,334, 375,151, 334,317, 251,626, 499,474, 750,326, 534,814, 1079,818 46 | anchors_p7_1280: 47 | - [19,22, 54,36, 32,77] # P3/8 48 | - [70,83, 138,71, 75,173] # P4/16 49 | - [165,159, 148,334, 375,151] # P5/32 50 | - [334,317, 251,626, 499,474] # P6/64 51 | - [750,326, 534,814, 1079,818] # P7/128 52 | 53 | # P7-1920: thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1920, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 29,34, 81,55, 47,115, 105,124, 207,107, 113,259, 247,238, 222,500, 563,227, 501,476, 376,939, 749,711, 1126,489, 801,1222, 1618,1227 54 | anchors_p7_1920: 55 | - [29,34, 81,55, 47,115] # P3/8 56 | - [105,124, 207,107, 113,259] # P4/16 57 | - [247,238, 222,500, 563,227] # P5/32 58 | - [501,476, 376,939, 749,711] # P6/64 59 | - [1126,489, 801,1222, 1618,1227] # P7/128 60 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3-SPP head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, SPP, [512, [5, 9, 13]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov3-tiny.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,14, 23,27, 37,58] # P4/16 9 | - [81,82, 135,169, 344,319] # P5/32 10 | 11 | # YOLOv3-tiny backbone 12 | backbone: 13 | # [from, number, module, args] 14 | [[-1, 1, Conv, [16, 3, 1]], # 0 15 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 1-P1/2 16 | [-1, 1, Conv, [32, 3, 1]], 17 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 3-P2/4 18 | [-1, 1, Conv, [64, 3, 1]], 19 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 5-P3/8 20 | [-1, 1, Conv, [128, 3, 1]], 21 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 7-P4/16 22 | [-1, 1, Conv, [256, 3, 1]], 23 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 9-P5/32 24 | [-1, 1, Conv, [512, 3, 1]], 25 | [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]], # 11 26 | [-1, 1, nn.MaxPool2d, [2, 1, 0]], # 12 27 | ] 28 | 29 | # YOLOv3-tiny head 30 | head: 31 | [[-1, 1, Conv, [1024, 3, 1]], 32 | [-1, 1, Conv, [256, 1, 1]], 33 | [-1, 1, Conv, [512, 3, 1]], # 15 (P5/32-large) 34 | 35 | [-2, 1, Conv, [128, 1, 1]], 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 38 | [-1, 1, Conv, [256, 3, 1]], # 19 (P4/16-medium) 39 | 40 | [[19, 15], 1, Detect, [nc, anchors]], # Detect(P4, P5) 41 | ] 42 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov3.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3 head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, Conv, [512, [1, 1]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5-bifpn.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 BiFPN head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14, 6], 1, Concat, [1]], # cat P4 <--- BiFPN change 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5-fpn.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 FPN head 28 | head: 29 | [[-1, 3, C3, [1024, False]], # 10 (P5/32-large) 30 | 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 3, C3, [512, False]], # 14 (P4/16-medium) 35 | 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 3, C3, [256, False]], # 18 (P3/8-small) 40 | 41 | [[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 42 | ] 43 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5-p2.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # auto-anchor evolves 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 13 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 14 | [-1, 3, C3, [128]], 15 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 16 | [-1, 6, C3, [256]], 17 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 18 | [-1, 9, C3, [512]], 19 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 20 | [-1, 3, C3, [1024]], 21 | [-1, 1, SPPF, [1024, 5]], # 9 22 | ] 23 | 24 | # YOLOv5 v6.0 head 25 | head: 26 | [[-1, 1, Conv, [512, 1, 1]], 27 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 28 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 29 | [-1, 3, C3, [512, False]], # 13 30 | 31 | [-1, 1, Conv, [256, 1, 1]], 32 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 33 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 34 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 35 | 36 | [-1, 1, Conv, [128, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 2], 1, Concat, [1]], # cat backbone P2 39 | [-1, 1, C3, [128, False]], # 21 (P2/4-xsmall) 40 | 41 | [-1, 1, Conv, [128, 3, 2]], 42 | [[-1, 18], 1, Concat, [1]], # cat head P3 43 | [-1, 3, C3, [256, False]], # 24 (P3/8-small) 44 | 45 | [-1, 1, Conv, [256, 3, 2]], 46 | [[-1, 14], 1, Concat, [1]], # cat head P4 47 | [-1, 3, C3, [512, False]], # 27 (P4/16-medium) 48 | 49 | [-1, 1, Conv, [512, 3, 2]], 50 | [[-1, 10], 1, Concat, [1]], # cat head P5 51 | [-1, 3, C3, [1024, False]], # 30 (P5/32-large) 52 | 53 | [[21, 24, 27, 30], 1, Detect, [nc, anchors]], # Detect(P2, P3, P4, P5) 54 | ] 55 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5-p6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # auto-anchor 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 13 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 14 | [-1, 3, C3, [128]], 15 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 16 | [-1, 6, C3, [256]], 17 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 18 | [-1, 9, C3, [512]], 19 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 20 | [-1, 3, C3, [768]], 21 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 22 | [-1, 3, C3, [1024]], 23 | [-1, 1, SPPF, [1024, 5]], # 11 24 | ] 25 | 26 | # YOLOv5 v6.0 head 27 | head: 28 | [[-1, 1, Conv, [768, 1, 1]], 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 31 | [-1, 3, C3, [768, False]], # 15 32 | 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 35 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 36 | [-1, 3, C3, [512, False]], # 19 37 | 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 40 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 41 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 42 | 43 | [-1, 1, Conv, [256, 3, 2]], 44 | [[-1, 20], 1, Concat, [1]], # cat head P4 45 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 46 | 47 | [-1, 1, Conv, [512, 3, 2]], 48 | [[-1, 16], 1, Concat, [1]], # cat head P5 49 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 50 | 51 | [-1, 1, Conv, [768, 3, 2]], 52 | [[-1, 12], 1, Concat, [1]], # cat head P6 53 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 54 | 55 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 56 | ] 57 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5-p7.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # auto-anchor 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 13 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 14 | [-1, 3, C3, [128]], 15 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 16 | [-1, 6, C3, [256]], 17 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 18 | [-1, 9, C3, [512]], 19 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 20 | [-1, 3, C3, [768]], 21 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 22 | [-1, 3, C3, [1024]], 23 | [-1, 1, Conv, [1280, 3, 2]], # 11-P7/128 24 | [-1, 3, C3, [1280]], 25 | [-1, 1, SPPF, [1280, 5]], # 13 26 | ] 27 | 28 | # YOLOv5 head 29 | head: 30 | [[-1, 1, Conv, [1024, 1, 1]], 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 10], 1, Concat, [1]], # cat backbone P6 33 | [-1, 3, C3, [1024, False]], # 17 34 | 35 | [-1, 1, Conv, [768, 1, 1]], 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 38 | [-1, 3, C3, [768, False]], # 21 39 | 40 | [-1, 1, Conv, [512, 1, 1]], 41 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 42 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 43 | [-1, 3, C3, [512, False]], # 25 44 | 45 | [-1, 1, Conv, [256, 1, 1]], 46 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 47 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 48 | [-1, 3, C3, [256, False]], # 29 (P3/8-small) 49 | 50 | [-1, 1, Conv, [256, 3, 2]], 51 | [[-1, 26], 1, Concat, [1]], # cat head P4 52 | [-1, 3, C3, [512, False]], # 32 (P4/16-medium) 53 | 54 | [-1, 1, Conv, [512, 3, 2]], 55 | [[-1, 22], 1, Concat, [1]], # cat head P5 56 | [-1, 3, C3, [768, False]], # 35 (P5/32-large) 57 | 58 | [-1, 1, Conv, [768, 3, 2]], 59 | [[-1, 18], 1, Concat, [1]], # cat head P6 60 | [-1, 3, C3, [1024, False]], # 38 (P6/64-xlarge) 61 | 62 | [-1, 1, Conv, [1024, 3, 2]], 63 | [[-1, 14], 1, Concat, [1]], # cat head P7 64 | [-1, 3, C3, [1280, False]], # 41 (P7/128-xxlarge) 65 | 66 | [[29, 32, 35, 38, 41], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6, P7) 67 | ] 68 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5-panet.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 PANet head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5l6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5m6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5n6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5s-ghost.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, GhostConv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3Ghost, [128]], 18 | [-1, 1, GhostConv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3Ghost, [256]], 20 | [-1, 1, GhostConv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3Ghost, [512]], 22 | [-1, 1, GhostConv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3Ghost, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, GhostConv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3Ghost, [512, False]], # 13 33 | 34 | [-1, 1, GhostConv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3Ghost, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, GhostConv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3Ghost, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, GhostConv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3Ghost, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5s-transformer.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3TR, [1024]], # 9 <--- C3TR() Transformer module 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5s6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /yolov5_obb/models/hub/yolov5x6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /yolov5_obb/models/yolov5l.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 3 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/yolov5m.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 3 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/yolov5n.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/yolov5s.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 3 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/models/yolov5x.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5_obb/requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # Base ---------------------------------------- 4 | matplotlib>=3.2.2 5 | numpy>=1.18.5 6 | opencv-python>=4.5.4 7 | Pillow>=7.1.2 8 | PyYAML>=5.3.1 9 | requests>=2.23.0 10 | scipy>=1.4.1 11 | torch>=1.7.0 12 | torchvision>=0.8.1 13 | tqdm>=4.41.0 14 | 15 | # Logging ------------------------------------- 16 | tensorboard>=2.4.1 17 | # wandb 18 | 19 | # Plotting ------------------------------------ 20 | pandas>=1.1.4 21 | seaborn>=0.11.0 22 | 23 | # Export -------------------------------------- 24 | # coremltools>=4.1 # CoreML export 25 | # onnx>=1.9.0 # ONNX export 26 | # onnx-simplifier>=0.3.6 # ONNX simplifier 27 | # scikit-learn==0.19.2 # CoreML quantization 28 | # tensorflow>=2.4.1 # TFLite export 29 | # tensorflowjs>=3.9.0 # TF.js export 30 | # openvino-dev # OpenVINO export 31 | 32 | # Extras -------------------------------------- 33 | # albumentations>=1.0.3 34 | # Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172 35 | # pycocotools>=2.0 # COCO mAP 36 | # roboflow 37 | thop # FLOPs computation 38 | -------------------------------------------------------------------------------- /yolov5_obb/split_txt.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os 4 | import random 5 | 6 | # obb data split 7 | annfilepath=r'/split/labelTxt/' 8 | saveBasePath=r'split/' 9 | train_percent=0.95 10 | total_file = os.listdir(annfilepath) 11 | num=len(total_file) 12 | list=range(num) 13 | tr=int(num*train_percent) 14 | train=random.sample(list,tr) 15 | ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w') 16 | fval = open(os.path.join(saveBasePath,'val.txt'), 'w') 17 | for i in list: 18 | name=total_file[i].split('.')[0]+'\n' 19 | if i in train: 20 | ftrain.write(name) 21 | else: 22 | fval.write(name) 23 | ftrain.close() 24 | fval.close() 25 | print("train size",tr) 26 | print("valid size",num-tr) 27 | -------------------------------------------------------------------------------- /yolov5_obb/tools/TestJson2VocClassTxt.py: -------------------------------------------------------------------------------- 1 | """ 2 | Yolov5-obb检测结果Json 文件转Voc Class Txt 3 | --json_path 输入的json文件路径 4 | --save_path 输出文件夹路径 5 | """ 6 | 7 | import os 8 | import json 9 | from tqdm import tqdm 10 | import argparse 11 | import shutil 12 | 13 | parser = argparse.ArgumentParser() 14 | parser.add_argument('--json_path', default='runs/val/exp/last_predictions.json',type=str, help="input: coco format(json)") 15 | parser.add_argument('--save_path', default='runs/val/exp/last_predictions_Txt', type=str, help="specify where to save the output dir of labels") 16 | arg = parser.parse_args() 17 | 18 | # For DOTA-v2.0 19 | dotav2_classnames = [ 'plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 20 | 'tennis-court', 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 21 | 'swimming-pool', 'helicopter', 'container-crane', 'airport', 'helipad'] 22 | # For DOTA-v1.5 23 | dotav15_classnames = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 'tennis-court', 24 | 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter', 'container-crane'] 25 | # For DOTA-v1.0 26 | datav1_classnames = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 'tennis-court', 27 | 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter'] 28 | 29 | DOTA_CLASSES = dotav15_classnames 30 | if __name__ == '__main__': 31 | json_file = arg.json_path # COCO Object Instance 类型的标注 32 | ana_txt_save_path = arg.save_path # 保存的路径 33 | 34 | data = json.load(open(json_file, 'r')) 35 | if os.path.exists(ana_txt_save_path): 36 | shutil.rmtree(ana_txt_save_path) # delete output folderX 37 | os.makedirs(ana_txt_save_path) 38 | 39 | for data_dict in data: 40 | img_name = data_dict["file_name"] 41 | score = data_dict["score"] 42 | poly = data_dict["poly"] 43 | classname = DOTA_CLASSES[data_dict["category_id"]-1] # COCO's category_id start from 1, not 0 44 | 45 | lines = "%s %s %s %s %s %s %s %s %s %s\n" % (img_name, score, poly[0],poly[1],poly[2],poly[3],poly[4],poly[5],poly[6],poly[7]) 46 | with open(str(ana_txt_save_path + '/Task1_' + classname) + '.txt', 'a') as f: 47 | f.writelines(lines) 48 | pass 49 | print("Done!") -------------------------------------------------------------------------------- /yolov5_obb/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | utils/initialization 4 | """ 5 | 6 | 7 | def notebook_init(verbose=True): 8 | # Check system software and hardware 9 | print('Checking setup...') 10 | 11 | import os 12 | import shutil 13 | 14 | from utils.general import check_requirements, emojis, is_colab 15 | from utils.torch_utils import select_device # imports 16 | 17 | check_requirements(('psutil', 'IPython')) 18 | import psutil 19 | from IPython import display # to display images and clear console output 20 | 21 | if is_colab(): 22 | shutil.rmtree('/content/sample_data', ignore_errors=True) # remove colab /sample_data directory 23 | 24 | if verbose: 25 | # System info 26 | # gb = 1 / 1000 ** 3 # bytes to GB 27 | gib = 1 / 1024 ** 3 # bytes to GiB 28 | ram = psutil.virtual_memory().total 29 | total, used, free = shutil.disk_usage("/") 30 | display.clear_output() 31 | s = f'({os.cpu_count()} CPUs, {ram * gib:.1f} GB RAM, {(total - free) * gib:.1f}/{total * gib:.1f} GB disk)' 32 | else: 33 | s = '' 34 | 35 | select_device(newline=False) 36 | print(emojis(f'Setup complete ✅ {s}')) 37 | return display 38 | -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/activations.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/activations.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/augmentations.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/augmentations.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/autoanchor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/autoanchor.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/autobatch.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/autobatch.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/callbacks.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/callbacks.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/datasets.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/datasets.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/downloads.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/downloads.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/general.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/general.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/loss.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/plots.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/plots.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/rboxs_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/rboxs_utils.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/__pycache__/torch_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/__pycache__/torch_utils.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/activations.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Activation functions 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | 11 | # SiLU https://arxiv.org/pdf/1606.08415.pdf ---------------------------------------------------------------------------- 12 | class SiLU(nn.Module): # export-friendly version of nn.SiLU() 13 | @staticmethod 14 | def forward(x): 15 | return x * torch.sigmoid(x) 16 | 17 | 18 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish() 19 | @staticmethod 20 | def forward(x): 21 | # return x * F.hardsigmoid(x) # for TorchScript and CoreML 22 | return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for TorchScript, CoreML and ONNX 23 | 24 | 25 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 26 | class Mish(nn.Module): 27 | @staticmethod 28 | def forward(x): 29 | return x * F.softplus(x).tanh() 30 | 31 | 32 | class MemoryEfficientMish(nn.Module): 33 | class F(torch.autograd.Function): 34 | @staticmethod 35 | def forward(ctx, x): 36 | ctx.save_for_backward(x) 37 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 38 | 39 | @staticmethod 40 | def backward(ctx, grad_output): 41 | x = ctx.saved_tensors[0] 42 | sx = torch.sigmoid(x) 43 | fx = F.softplus(x).tanh() 44 | return grad_output * (fx + x * sx * (1 - fx * fx)) 45 | 46 | def forward(self, x): 47 | return self.F.apply(x) 48 | 49 | 50 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 51 | class FReLU(nn.Module): 52 | def __init__(self, c1, k=3): # ch_in, kernel 53 | super().__init__() 54 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False) 55 | self.bn = nn.BatchNorm2d(c1) 56 | 57 | def forward(self, x): 58 | return torch.max(x, self.bn(self.conv(x))) 59 | 60 | 61 | # ACON https://arxiv.org/pdf/2009.04759.pdf ---------------------------------------------------------------------------- 62 | class AconC(nn.Module): 63 | r""" ACON activation (activate or not). 64 | AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter 65 | according to "Activate or Not: Learning Customized Activation" . 66 | """ 67 | 68 | def __init__(self, c1): 69 | super().__init__() 70 | self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) 71 | self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1)) 72 | self.beta = nn.Parameter(torch.ones(1, c1, 1, 1)) 73 | 74 | def forward(self, x): 75 | dpx = (self.p1 - self.p2) * x 76 | return dpx * torch.sigmoid(self.beta * dpx) + self.p2 * x 77 | 78 | 79 | class MetaAconC(nn.Module): 80 | r""" ACON activation (activate or not). 81 | MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network 82 | according to "Activate or Not: Learning Customized Activation" . 83 | """ 84 | 85 | def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r 86 | super().__init__() 87 | c2 = max(r, c1 // r) 88 | self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) 89 | self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1)) 90 | self.fc1 = nn.Conv2d(c1, c2, k, s, bias=True) 91 | self.fc2 = nn.Conv2d(c2, c1, k, s, bias=True) 92 | # self.bn1 = nn.BatchNorm2d(c2) 93 | # self.bn2 = nn.BatchNorm2d(c1) 94 | 95 | def forward(self, x): 96 | y = x.mean(dim=2, keepdims=True).mean(dim=3, keepdims=True) 97 | # batch-size 1 bug/instabilities https://github.com/ultralytics/yolov5/issues/2891 98 | # beta = torch.sigmoid(self.bn2(self.fc2(self.bn1(self.fc1(y))))) # bug/unstable 99 | beta = torch.sigmoid(self.fc2(self.fc1(y))) # bug patch BN layers removed 100 | dpx = (self.p1 - self.p2) * x 101 | return dpx * torch.sigmoid(beta * dpx) + self.p2 * x 102 | -------------------------------------------------------------------------------- /yolov5_obb/utils/autobatch.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Auto-batch utils 4 | """ 5 | 6 | from copy import deepcopy 7 | 8 | import numpy as np 9 | import torch 10 | from torch.cuda import amp 11 | 12 | from utils.general import LOGGER, colorstr 13 | from utils.torch_utils import profile 14 | 15 | 16 | def check_train_batch_size(model, imgsz=640): 17 | # Check YOLOv5 training batch size 18 | with amp.autocast(): 19 | return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size 20 | 21 | 22 | def autobatch(model, imgsz=640, fraction=0.9, batch_size=16): 23 | # Automatically estimate best batch size to use `fraction` of available CUDA memory 24 | # Usage: 25 | # import torch 26 | # from utils.autobatch import autobatch 27 | # model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False) 28 | # print(autobatch(model)) 29 | 30 | prefix = colorstr('AutoBatch: ') 31 | LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}') 32 | device = next(model.parameters()).device # get model device 33 | if device.type == 'cpu': 34 | LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}') 35 | return batch_size 36 | 37 | d = str(device).upper() # 'CUDA:0' 38 | properties = torch.cuda.get_device_properties(device) # device properties 39 | t = properties.total_memory / 1024 ** 3 # (GiB) 40 | r = torch.cuda.memory_reserved(device) / 1024 ** 3 # (GiB) 41 | a = torch.cuda.memory_allocated(device) / 1024 ** 3 # (GiB) 42 | f = t - (r + a) # free inside reserved 43 | LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free') 44 | 45 | batch_sizes = [1, 2, 4, 8, 16] 46 | try: 47 | img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes] 48 | y = profile(img, model, n=3, device=device) 49 | except Exception as e: 50 | LOGGER.warning(f'{prefix}{e}') 51 | 52 | y = [x[2] for x in y if x] # memory [2] 53 | batch_sizes = batch_sizes[:len(y)] 54 | p = np.polyfit(batch_sizes, y, deg=1) # first degree polynomial fit 55 | b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size) 56 | LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)') 57 | return b 58 | -------------------------------------------------------------------------------- /yolov5_obb/utils/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/aws/__init__.py -------------------------------------------------------------------------------- /yolov5_obb/utils/aws/mime.sh: -------------------------------------------------------------------------------- 1 | # AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/ 2 | # This script will run on every instance restart, not only on first start 3 | # --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA --- 4 | 5 | Content-Type: multipart/mixed; boundary="//" 6 | MIME-Version: 1.0 7 | 8 | --// 9 | Content-Type: text/cloud-config; charset="us-ascii" 10 | MIME-Version: 1.0 11 | Content-Transfer-Encoding: 7bit 12 | Content-Disposition: attachment; filename="cloud-config.txt" 13 | 14 | #cloud-config 15 | cloud_final_modules: 16 | - [scripts-user, always] 17 | 18 | --// 19 | Content-Type: text/x-shellscript; charset="us-ascii" 20 | MIME-Version: 1.0 21 | Content-Transfer-Encoding: 7bit 22 | Content-Disposition: attachment; filename="userdata.txt" 23 | 24 | #!/bin/bash 25 | # --- paste contents of userdata.sh here --- 26 | --// 27 | -------------------------------------------------------------------------------- /yolov5_obb/utils/aws/resume.py: -------------------------------------------------------------------------------- 1 | # Resume all interrupted trainings in yolov5/ dir including DDP trainings 2 | # Usage: $ python utils/aws/resume.py 3 | 4 | import os 5 | import sys 6 | from pathlib import Path 7 | 8 | import torch 9 | import yaml 10 | 11 | FILE = Path(__file__).resolve() 12 | ROOT = FILE.parents[2] # YOLOv5 root directory 13 | if str(ROOT) not in sys.path: 14 | sys.path.append(str(ROOT)) # add ROOT to PATH 15 | 16 | port = 0 # --master_port 17 | path = Path('').resolve() 18 | for last in path.rglob('*/**/last.pt'): 19 | ckpt = torch.load(last) 20 | if ckpt['optimizer'] is None: 21 | continue 22 | 23 | # Load opt.yaml 24 | with open(last.parent.parent / 'opt.yaml', errors='ignore') as f: 25 | opt = yaml.safe_load(f) 26 | 27 | # Get device count 28 | d = opt['device'].split(',') # devices 29 | nd = len(d) # number of devices 30 | ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel 31 | 32 | if ddp: # multi-GPU 33 | port += 1 34 | cmd = f'python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}' 35 | else: # single-GPU 36 | cmd = f'python train.py --resume {last}' 37 | 38 | cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread 39 | print(cmd) 40 | os.system(cmd) 41 | -------------------------------------------------------------------------------- /yolov5_obb/utils/aws/userdata.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html 3 | # This script will run only once on first instance start (for a re-start script see mime.sh) 4 | # /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir 5 | # Use >300 GB SSD 6 | 7 | cd home/ubuntu 8 | if [ ! -d yolov5 ]; then 9 | echo "Running first-time script." # install dependencies, download COCO, pull Docker 10 | git clone https://github.com/ultralytics/yolov5 -b master && sudo chmod -R 777 yolov5 11 | cd yolov5 12 | bash data/scripts/get_coco.sh && echo "COCO done." & 13 | sudo docker pull ultralytics/yolov5:latest && echo "Docker done." & 14 | python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." & 15 | wait && echo "All tasks done." # finish background tasks 16 | else 17 | echo "Running re-start script." # resume interrupted runs 18 | i=0 19 | list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour' 20 | while IFS= read -r id; do 21 | ((i++)) 22 | echo "restarting container $i: $id" 23 | sudo docker start $id 24 | # sudo docker exec -it $id python train.py --resume # single-GPU 25 | sudo docker exec -d $id python utils/aws/resume.py # multi-scenario 26 | done <<<"$list" 27 | fi 28 | -------------------------------------------------------------------------------- /yolov5_obb/utils/callbacks.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Callback utils 4 | """ 5 | 6 | 7 | class Callbacks: 8 | """" 9 | Handles all registered callbacks for YOLOv5 Hooks 10 | """ 11 | 12 | def __init__(self): 13 | # Define the available callbacks 14 | self._callbacks = { 15 | 'on_pretrain_routine_start': [], 16 | 'on_pretrain_routine_end': [], 17 | 18 | 'on_train_start': [], 19 | 'on_train_epoch_start': [], 20 | 'on_train_batch_start': [], 21 | 'optimizer_step': [], 22 | 'on_before_zero_grad': [], 23 | 'on_train_batch_end': [], 24 | 'on_train_epoch_end': [], 25 | 26 | 'on_val_start': [], 27 | 'on_val_batch_start': [], 28 | 'on_val_image_end': [], 29 | 'on_val_batch_end': [], 30 | 'on_val_end': [], 31 | 32 | 'on_fit_epoch_end': [], # fit = train + val 33 | 'on_model_save': [], 34 | 'on_train_end': [], 35 | 'on_params_update': [], 36 | 'teardown': [], 37 | } 38 | 39 | def register_action(self, hook, name='', callback=None): 40 | """ 41 | Register a new action to a callback hook 42 | 43 | Args: 44 | hook The callback hook name to register the action to 45 | name The name of the action for later reference 46 | callback The callback to fire 47 | """ 48 | assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" 49 | assert callable(callback), f"callback '{callback}' is not callable" 50 | self._callbacks[hook].append({'name': name, 'callback': callback}) 51 | 52 | def get_registered_actions(self, hook=None): 53 | """" 54 | Returns all the registered actions by callback hook 55 | 56 | Args: 57 | hook The name of the hook to check, defaults to all 58 | """ 59 | if hook: 60 | return self._callbacks[hook] 61 | else: 62 | return self._callbacks 63 | 64 | def run(self, hook, *args, **kwargs): 65 | """ 66 | Loop through the registered actions and fire all callbacks 67 | 68 | Args: 69 | hook The name of the hook to check, defaults to all 70 | args Arguments to receive from YOLOv5 71 | kwargs Keyword Arguments to receive from YOLOv5 72 | """ 73 | 74 | assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" 75 | 76 | for logger in self._callbacks[hook]: 77 | logger['callback'](*args, **kwargs) 78 | -------------------------------------------------------------------------------- /yolov5_obb/utils/downloads.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Download utils 4 | """ 5 | 6 | import os 7 | import platform 8 | import subprocess 9 | import time 10 | import urllib 11 | from pathlib import Path 12 | from zipfile import ZipFile 13 | 14 | import requests 15 | import torch 16 | 17 | 18 | def gsutil_getsize(url=''): 19 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 20 | s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8') 21 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 22 | 23 | 24 | def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''): 25 | # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes 26 | file = Path(file) 27 | assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}" 28 | try: # url1 29 | print(f'Downloading {url} to {file}...') 30 | torch.hub.download_url_to_file(url, str(file)) 31 | assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check 32 | except Exception as e: # url2 33 | file.unlink(missing_ok=True) # remove partial downloads 34 | print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...') 35 | os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail 36 | finally: 37 | if not file.exists() or file.stat().st_size < min_bytes: # check 38 | file.unlink(missing_ok=True) # remove partial downloads 39 | print(f"ERROR: {assert_msg}\n{error_msg}") 40 | print('') 41 | 42 | 43 | def attempt_download(file, repo='ultralytics/yolov5'): # from utils.downloads import *; attempt_download() 44 | # Attempt file download if does not exist 45 | file = Path(str(file).strip().replace("'", '')) 46 | 47 | if not file.exists(): 48 | # URL specified 49 | name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc. 50 | if str(file).startswith(('http:/', 'https:/')): # download 51 | url = str(file).replace(':/', '://') # Pathlib turns :// -> :/ 52 | file = name.split('?')[0] # parse authentication https://url.com/file.txt?auth... 53 | if Path(file).is_file(): 54 | print(f'Found {url} locally at {file}') # file already exists 55 | else: 56 | safe_download(file=file, url=url, min_bytes=1E5) 57 | return file 58 | 59 | # GitHub assets 60 | file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required) 61 | try: 62 | response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api 63 | assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] 64 | tag = response['tag_name'] # i.e. 'v1.0' 65 | except: # fallback plan 66 | assets = ['yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 67 | 'yolov5n6.pt', 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt'] 68 | try: 69 | tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1] 70 | except: 71 | tag = 'v6.0' # current release 72 | 73 | if name in assets: 74 | safe_download(file, 75 | url=f'https://github.com/{repo}/releases/download/{tag}/{name}', 76 | # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional) 77 | min_bytes=1E5, 78 | error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/') 79 | 80 | return str(file) 81 | 82 | 83 | def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'): 84 | # Downloads a file from Google Drive. from yolov5.utils.downloads import *; gdrive_download() 85 | t = time.time() 86 | file = Path(file) 87 | cookie = Path('cookie') # gdrive cookie 88 | print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='') 89 | file.unlink(missing_ok=True) # remove existing file 90 | cookie.unlink(missing_ok=True) # remove existing cookie 91 | 92 | # Attempt file download 93 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 94 | os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}') 95 | if os.path.exists('cookie'): # large file 96 | s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}' 97 | else: # small file 98 | s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"' 99 | r = os.system(s) # execute, capture return 100 | cookie.unlink(missing_ok=True) # remove existing cookie 101 | 102 | # Error check 103 | if r != 0: 104 | file.unlink(missing_ok=True) # remove partial 105 | print('Download error ') # raise Exception('Download error') 106 | return r 107 | 108 | # Unzip if archive 109 | if file.suffix == '.zip': 110 | print('unzipping... ', end='') 111 | ZipFile(file).extractall(path=file.parent) # unzip 112 | file.unlink() # remove zip 113 | 114 | print(f'Done ({time.time() - t:.1f}s)') 115 | return r 116 | 117 | 118 | def get_token(cookie="./cookie"): 119 | with open(cookie) as f: 120 | for line in f: 121 | if "download" in line: 122 | return line.split()[-1] 123 | return "" 124 | 125 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries ---------------------------------------------- 126 | # 127 | # 128 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 129 | # # Uploads a file to a bucket 130 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 131 | # 132 | # storage_client = storage.Client() 133 | # bucket = storage_client.get_bucket(bucket_name) 134 | # blob = bucket.blob(destination_blob_name) 135 | # 136 | # blob.upload_from_filename(source_file_name) 137 | # 138 | # print('File {} uploaded to {}.'.format( 139 | # source_file_name, 140 | # destination_blob_name)) 141 | # 142 | # 143 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 144 | # # Uploads a blob from a bucket 145 | # storage_client = storage.Client() 146 | # bucket = storage_client.get_bucket(bucket_name) 147 | # blob = bucket.blob(source_blob_name) 148 | # 149 | # blob.download_to_filename(destination_file_name) 150 | # 151 | # print('Blob {} downloaded to {}.'.format( 152 | # source_blob_name, 153 | # destination_file_name)) 154 | -------------------------------------------------------------------------------- /yolov5_obb/utils/flask_rest_api/README.md: -------------------------------------------------------------------------------- 1 | # Flask REST API 2 | 3 | [REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are 4 | commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API 5 | created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/). 6 | 7 | ## Requirements 8 | 9 | [Flask](https://palletsprojects.com/p/flask/) is required. Install with: 10 | 11 | ```shell 12 | $ pip install Flask 13 | ``` 14 | 15 | ## Run 16 | 17 | After Flask installation run: 18 | 19 | ```shell 20 | $ python3 restapi.py --port 5000 21 | ``` 22 | 23 | Then use [curl](https://curl.se/) to perform a request: 24 | 25 | ```shell 26 | $ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s' 27 | ``` 28 | 29 | The model inference results are returned as a JSON response: 30 | 31 | ```json 32 | [ 33 | { 34 | "class": 0, 35 | "confidence": 0.8900438547, 36 | "height": 0.9318675399, 37 | "name": "person", 38 | "width": 0.3264600933, 39 | "xcenter": 0.7438579798, 40 | "ycenter": 0.5207948685 41 | }, 42 | { 43 | "class": 0, 44 | "confidence": 0.8440024257, 45 | "height": 0.7155083418, 46 | "name": "person", 47 | "width": 0.6546785235, 48 | "xcenter": 0.427829951, 49 | "ycenter": 0.6334488392 50 | }, 51 | { 52 | "class": 27, 53 | "confidence": 0.3771208823, 54 | "height": 0.3902671337, 55 | "name": "tie", 56 | "width": 0.0696444362, 57 | "xcenter": 0.3675483763, 58 | "ycenter": 0.7991207838 59 | }, 60 | { 61 | "class": 27, 62 | "confidence": 0.3527112305, 63 | "height": 0.1540903747, 64 | "name": "tie", 65 | "width": 0.0336618312, 66 | "xcenter": 0.7814827561, 67 | "ycenter": 0.5065554976 68 | } 69 | ] 70 | ``` 71 | 72 | An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given 73 | in `example_request.py` 74 | -------------------------------------------------------------------------------- /yolov5_obb/utils/flask_rest_api/example_request.py: -------------------------------------------------------------------------------- 1 | """Perform test request""" 2 | import pprint 3 | 4 | import requests 5 | 6 | DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s" 7 | TEST_IMAGE = "zidane.jpg" 8 | 9 | image_data = open(TEST_IMAGE, "rb").read() 10 | 11 | response = requests.post(DETECTION_URL, files={"image": image_data}).json() 12 | 13 | pprint.pprint(response) 14 | -------------------------------------------------------------------------------- /yolov5_obb/utils/flask_rest_api/restapi.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run a rest API exposing the yolov5s object detection model 3 | """ 4 | import argparse 5 | import io 6 | 7 | import torch 8 | from flask import Flask, request 9 | from PIL import Image 10 | 11 | app = Flask(__name__) 12 | 13 | DETECTION_URL = "/v1/object-detection/yolov5s" 14 | 15 | 16 | @app.route(DETECTION_URL, methods=["POST"]) 17 | def predict(): 18 | if not request.method == "POST": 19 | return 20 | 21 | if request.files.get("image"): 22 | image_file = request.files["image"] 23 | image_bytes = image_file.read() 24 | 25 | img = Image.open(io.BytesIO(image_bytes)) 26 | 27 | results = model(img, size=640) # reduce size=320 for faster inference 28 | return results.pandas().xyxy[0].to_json(orient="records") 29 | 30 | 31 | if __name__ == "__main__": 32 | parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model") 33 | parser.add_argument("--port", default=5000, type=int, help="port number") 34 | args = parser.parse_args() 35 | 36 | model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force_reload to recache 37 | app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat 38 | -------------------------------------------------------------------------------- /yolov5_obb/utils/google_app_engine/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM gcr.io/google-appengine/python 2 | 3 | # Create a virtualenv for dependencies. This isolates these packages from 4 | # system-level packages. 5 | # Use -p python3 or -p python3.7 to select python version. Default is version 2. 6 | RUN virtualenv /env -p python3 7 | 8 | # Setting these environment variables are the same as running 9 | # source /env/bin/activate. 10 | ENV VIRTUAL_ENV /env 11 | ENV PATH /env/bin:$PATH 12 | 13 | RUN apt-get update && apt-get install -y python-opencv 14 | 15 | # Copy the application's requirements.txt and run pip to install all 16 | # dependencies into the virtualenv. 17 | ADD requirements.txt /app/requirements.txt 18 | RUN pip install -r /app/requirements.txt 19 | 20 | # Add the application source code. 21 | ADD . /app 22 | 23 | # Run a WSGI server to serve the application. gunicorn must be declared as 24 | # a dependency in requirements.txt. 25 | CMD gunicorn -b :$PORT main:app 26 | -------------------------------------------------------------------------------- /yolov5_obb/utils/google_app_engine/additional_requirements.txt: -------------------------------------------------------------------------------- 1 | # add these requirements in your app on top of the existing ones 2 | pip==21.1 3 | Flask==1.0.2 4 | gunicorn==19.9.0 5 | -------------------------------------------------------------------------------- /yolov5_obb/utils/google_app_engine/app.yaml: -------------------------------------------------------------------------------- 1 | runtime: custom 2 | env: flex 3 | 4 | service: yolov5app 5 | 6 | liveness_check: 7 | initial_delay_sec: 600 8 | 9 | manual_scaling: 10 | instances: 1 11 | resources: 12 | cpu: 1 13 | memory_gb: 4 14 | disk_size_gb: 20 15 | -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/wandb/__init__.py -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/wandb/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/wandb/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/__pycache__/wandb_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/wandb/__pycache__/wandb_utils.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/__pycache__/wandb_utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/loggers/wandb/__pycache__/wandb_utils.cpython-39.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/log_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from wandb_utils import WandbLogger 4 | 5 | from utils.general import LOGGER 6 | 7 | WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' 8 | 9 | 10 | def create_dataset_artifact(opt): 11 | logger = WandbLogger(opt, None, job_type='Dataset Creation') # TODO: return value unused 12 | if not logger.wandb: 13 | LOGGER.info("install wandb using `pip install wandb` to log the dataset") 14 | 15 | 16 | if __name__ == '__main__': 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path') 19 | parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') 20 | parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project') 21 | parser.add_argument('--entity', default=None, help='W&B entity') 22 | parser.add_argument('--name', type=str, default='log dataset', help='name of W&B run') 23 | 24 | opt = parser.parse_args() 25 | opt.resume = False # Explicitly disallow resume check for dataset upload job 26 | 27 | create_dataset_artifact(opt) 28 | -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/sweep.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pathlib import Path 3 | 4 | import wandb 5 | 6 | FILE = Path(__file__).resolve() 7 | ROOT = FILE.parents[3] # YOLOv5 root directory 8 | if str(ROOT) not in sys.path: 9 | sys.path.append(str(ROOT)) # add ROOT to PATH 10 | 11 | from train import parse_opt, train 12 | from utils.callbacks import Callbacks 13 | from utils.general import increment_path 14 | from utils.torch_utils import select_device 15 | 16 | 17 | def sweep(): 18 | wandb.init() 19 | # Get hyp dict from sweep agent 20 | hyp_dict = vars(wandb.config).get("_items") 21 | 22 | # Workaround: get necessary opt args 23 | opt = parse_opt(known=True) 24 | opt.batch_size = hyp_dict.get("batch_size") 25 | opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve)) 26 | opt.epochs = hyp_dict.get("epochs") 27 | opt.nosave = True 28 | opt.data = hyp_dict.get("data") 29 | opt.weights = str(opt.weights) 30 | opt.cfg = str(opt.cfg) 31 | opt.data = str(opt.data) 32 | opt.hyp = str(opt.hyp) 33 | opt.project = str(opt.project) 34 | device = select_device(opt.device, batch_size=opt.batch_size) 35 | 36 | # train 37 | train(hyp_dict, opt, device, callbacks=Callbacks()) 38 | 39 | 40 | if __name__ == "__main__": 41 | sweep() 42 | -------------------------------------------------------------------------------- /yolov5_obb/utils/loggers/wandb/sweep.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for training 2 | # To set range- 3 | # Provide min and max values as: 4 | # parameter: 5 | # 6 | # min: scalar 7 | # max: scalar 8 | # OR 9 | # 10 | # Set a specific list of search space- 11 | # parameter: 12 | # values: [scalar1, scalar2, scalar3...] 13 | # 14 | # You can use grid, bayesian and hyperopt search strategy 15 | # For more info on configuring sweeps visit - https://docs.wandb.ai/guides/sweeps/configuration 16 | 17 | program: utils/loggers/wandb/sweep.py 18 | method: random 19 | metric: 20 | name: metrics/mAP_0.5 21 | goal: maximize 22 | 23 | parameters: 24 | # hyperparameters: set either min, max range or values list 25 | data: 26 | value: "data/coco128.yaml" 27 | batch_size: 28 | values: [64] 29 | epochs: 30 | values: [10] 31 | 32 | lr0: 33 | distribution: uniform 34 | min: 1e-5 35 | max: 1e-1 36 | lrf: 37 | distribution: uniform 38 | min: 0.01 39 | max: 1.0 40 | momentum: 41 | distribution: uniform 42 | min: 0.6 43 | max: 0.98 44 | weight_decay: 45 | distribution: uniform 46 | min: 0.0 47 | max: 0.001 48 | warmup_epochs: 49 | distribution: uniform 50 | min: 0.0 51 | max: 5.0 52 | warmup_momentum: 53 | distribution: uniform 54 | min: 0.0 55 | max: 0.95 56 | warmup_bias_lr: 57 | distribution: uniform 58 | min: 0.0 59 | max: 0.2 60 | box: 61 | distribution: uniform 62 | min: 0.02 63 | max: 0.2 64 | cls: 65 | distribution: uniform 66 | min: 0.2 67 | max: 4.0 68 | cls_pw: 69 | distribution: uniform 70 | min: 0.5 71 | max: 2.0 72 | obj: 73 | distribution: uniform 74 | min: 0.2 75 | max: 4.0 76 | obj_pw: 77 | distribution: uniform 78 | min: 0.5 79 | max: 2.0 80 | iou_t: 81 | distribution: uniform 82 | min: 0.1 83 | max: 0.7 84 | anchor_t: 85 | distribution: uniform 86 | min: 2.0 87 | max: 8.0 88 | fl_gamma: 89 | distribution: uniform 90 | min: 0.0 91 | max: 0.1 92 | hsv_h: 93 | distribution: uniform 94 | min: 0.0 95 | max: 0.1 96 | hsv_s: 97 | distribution: uniform 98 | min: 0.0 99 | max: 0.9 100 | hsv_v: 101 | distribution: uniform 102 | min: 0.0 103 | max: 0.9 104 | degrees: 105 | distribution: uniform 106 | min: 0.0 107 | max: 45.0 108 | translate: 109 | distribution: uniform 110 | min: 0.0 111 | max: 0.9 112 | scale: 113 | distribution: uniform 114 | min: 0.0 115 | max: 0.9 116 | shear: 117 | distribution: uniform 118 | min: 0.0 119 | max: 10.0 120 | perspective: 121 | distribution: uniform 122 | min: 0.0 123 | max: 0.001 124 | flipud: 125 | distribution: uniform 126 | min: 0.0 127 | max: 1.0 128 | fliplr: 129 | distribution: uniform 130 | min: 0.0 131 | max: 1.0 132 | mosaic: 133 | distribution: uniform 134 | min: 0.0 135 | max: 1.0 136 | mixup: 137 | distribution: uniform 138 | min: 0.0 139 | max: 1.0 140 | copy_paste: 141 | distribution: uniform 142 | min: 0.0 143 | max: 1.0 144 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/__init__.py: -------------------------------------------------------------------------------- 1 | from .nms_rotated_wrapper import obb_nms, poly_nms 2 | 3 | __all__ = ['obb_nms', 'poly_nms'] 4 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/__pycache__/nms_rotated_wrapper.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/__pycache__/nms_rotated_wrapper.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/build/lib.linux-x86_64-3.7/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/build/lib.linux-x86_64-3.7/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/nms_rotated_cpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/nms_rotated_cpu.o -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/nms_rotated_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/nms_rotated_cuda.o -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/nms_rotated_ext.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/nms_rotated_ext.o -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/poly_nms_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/build/temp.linux-x86_64-3.7/src/poly_nms_cuda.o -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/utils/nms_rotated/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_ext.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: nms-rotated-ext 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_ext.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | setup.py 2 | nms_rotated_ext.egg-info/PKG-INFO 3 | nms_rotated_ext.egg-info/SOURCES.txt 4 | nms_rotated_ext.egg-info/dependency_links.txt 5 | nms_rotated_ext.egg-info/not-zip-safe 6 | nms_rotated_ext.egg-info/top_level.txt 7 | src/nms_rotated_cpu.cpp 8 | src/nms_rotated_cuda.cu 9 | src/nms_rotated_ext.cpp 10 | src/poly_nms_cuda.cu -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_ext.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_ext.egg-info/not-zip-safe: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_ext.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/nms_rotated_wrapper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from . import nms_rotated_ext 5 | 6 | def obb_nms(dets, scores, iou_thr, device_id=None): 7 | """ 8 | RIoU NMS - iou_thr. 9 | Args: 10 | dets (tensor/array): (num, [cx cy w h θ]) θ∈[-pi/2, pi/2) 11 | scores (tensor/array): (num) 12 | iou_thr (float): (1) 13 | Returns: 14 | dets (tensor): (n_nms, [cx cy w h θ]) 15 | inds (tensor): (n_nms), nms index of dets 16 | """ 17 | if isinstance(dets, torch.Tensor): 18 | is_numpy = False 19 | dets_th = dets 20 | elif isinstance(dets, np.ndarray): 21 | is_numpy = True 22 | device = 'cpu' if device_id is None else f'cuda:{device_id}' 23 | dets_th = torch.from_numpy(dets).to(device) 24 | else: 25 | raise TypeError('dets must be eithr a Tensor or numpy array, ' 26 | f'but got {type(dets)}') 27 | 28 | if dets_th.numel() == 0: # len(dets) 29 | inds = dets_th.new_zeros(0, dtype=torch.int64) 30 | else: 31 | # same bug will happen when bboxes is too small 32 | too_small = dets_th[:, [2, 3]].min(1)[0] < 0.001 # [n] 33 | if too_small.all(): # all the bboxes is too small 34 | inds = dets_th.new_zeros(0, dtype=torch.int64) 35 | else: 36 | ori_inds = torch.arange(dets_th.size(0)) # 0 ~ n-1 37 | ori_inds = ori_inds[~too_small] 38 | dets_th = dets_th[~too_small] # (n_filter, 5) 39 | scores = scores[~too_small] 40 | 41 | inds = nms_rotated_ext.nms_rotated(dets_th, scores, iou_thr) 42 | inds = ori_inds[inds] 43 | 44 | if is_numpy: 45 | inds = inds.cpu().numpy() 46 | return dets[inds, :], inds 47 | 48 | 49 | def poly_nms(dets, iou_thr, device_id=None): 50 | if isinstance(dets, torch.Tensor): 51 | is_numpy = False 52 | dets_th = dets 53 | elif isinstance(dets, np.ndarray): 54 | is_numpy = True 55 | device = 'cpu' if device_id is None else f'cuda:{device_id}' 56 | dets_th = torch.from_numpy(dets).to(device) 57 | else: 58 | raise TypeError('dets must be eithr a Tensor or numpy array, ' 59 | f'but got {type(dets)}') 60 | 61 | if dets_th.device == torch.device('cpu'): 62 | raise NotImplementedError 63 | inds = nms_rotated_ext.nms_poly(dets_th.float(), iou_thr) 64 | 65 | if is_numpy: 66 | inds = inds.cpu().numpy() 67 | return dets[inds, :], inds 68 | 69 | if __name__ == '__main__': 70 | rboxes_opencv = torch.tensor(([136.6, 111.6, 200, 100, -60], 71 | [136.6, 111.6, 100, 200, -30], 72 | [100, 100, 141.4, 141.4, -45], 73 | [100, 100, 141.4, 141.4, -45])) 74 | rboxes_longedge = torch.tensor(([136.6, 111.6, 200, 100, -60], 75 | [136.6, 111.6, 200, 100, 120], 76 | [100, 100, 141.4, 141.4, 45], 77 | [100, 100, 141.4, 141.4, 135])) 78 | 79 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/src/nms_rotated_cpu.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/facebookresearch/detectron2/tree/master/detectron2/layers/csrc/nms_rotated 3 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 4 | #include 5 | #include "box_iou_rotated_utils.h" 6 | 7 | 8 | template 9 | at::Tensor nms_rotated_cpu_kernel( 10 | const at::Tensor& dets, 11 | const at::Tensor& scores, 12 | const float iou_threshold) { 13 | // nms_rotated_cpu_kernel is modified from torchvision's nms_cpu_kernel, 14 | // however, the code in this function is much shorter because 15 | // we delegate the IoU computation for rotated boxes to 16 | // the single_box_iou_rotated function in box_iou_rotated_utils.h 17 | AT_ASSERTM(dets.device().is_cpu(), "dets must be a CPU tensor"); 18 | AT_ASSERTM(scores.device().is_cpu(), "scores must be a CPU tensor"); 19 | AT_ASSERTM( 20 | dets.scalar_type() == scores.scalar_type(), 21 | "dets should have the same type as scores"); 22 | 23 | if (dets.numel() == 0) { 24 | return at::empty({0}, dets.options().dtype(at::kLong)); 25 | } 26 | 27 | auto order_t = std::get<1>(scores.sort(0, /* descending=*/true)); 28 | 29 | auto ndets = dets.size(0); 30 | at::Tensor suppressed_t = at::zeros({ndets}, dets.options().dtype(at::kByte)); 31 | at::Tensor keep_t = at::zeros({ndets}, dets.options().dtype(at::kLong)); 32 | 33 | auto suppressed = suppressed_t.data_ptr(); 34 | auto keep = keep_t.data_ptr(); 35 | auto order = order_t.data_ptr(); 36 | 37 | int64_t num_to_keep = 0; 38 | 39 | for (int64_t _i = 0; _i < ndets; _i++) { 40 | auto i = order[_i]; 41 | if (suppressed[i] == 1) { 42 | continue; 43 | } 44 | 45 | keep[num_to_keep++] = i; 46 | 47 | for (int64_t _j = _i + 1; _j < ndets; _j++) { 48 | auto j = order[_j]; 49 | if (suppressed[j] == 1) { 50 | continue; 51 | } 52 | 53 | auto ovr = single_box_iou_rotated( 54 | dets[i].data_ptr(), dets[j].data_ptr()); 55 | if (ovr >= iou_threshold) { 56 | suppressed[j] = 1; 57 | } 58 | } 59 | } 60 | return keep_t.narrow(/*dim=*/0, /*start=*/0, /*length=*/num_to_keep); 61 | } 62 | 63 | at::Tensor nms_rotated_cpu( 64 | // input must be contiguous 65 | const at::Tensor& dets, 66 | const at::Tensor& scores, 67 | const float iou_threshold) { 68 | auto result = at::empty({0}, dets.options()); 69 | 70 | AT_DISPATCH_FLOATING_TYPES(dets.scalar_type(), "nms_rotated", [&] { 71 | result = nms_rotated_cpu_kernel(dets, scores, iou_threshold); 72 | }); 73 | return result; 74 | } 75 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/src/nms_rotated_cuda.cu: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/facebookresearch/detectron2/tree/master/detectron2/layers/csrc/nms_rotated 3 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "box_iou_rotated_utils.h" 9 | 10 | int const threadsPerBlock = sizeof(unsigned long long) * 8; 11 | 12 | template 13 | __global__ void nms_rotated_cuda_kernel( 14 | const int n_boxes, 15 | const float iou_threshold, 16 | const T* dev_boxes, 17 | unsigned long long* dev_mask) { 18 | // nms_rotated_cuda_kernel is modified from torchvision's nms_cuda_kernel 19 | 20 | const int row_start = blockIdx.y; 21 | const int col_start = blockIdx.x; 22 | 23 | // if (row_start > col_start) return; 24 | 25 | const int row_size = 26 | min(n_boxes - row_start * threadsPerBlock, threadsPerBlock); 27 | const int col_size = 28 | min(n_boxes - col_start * threadsPerBlock, threadsPerBlock); 29 | 30 | // Compared to nms_cuda_kernel, where each box is represented with 4 values 31 | // (x1, y1, x2, y2), each rotated box is represented with 5 values 32 | // (x_center, y_center, width, height, angle_degrees) here. 33 | __shared__ T block_boxes[threadsPerBlock * 5]; 34 | if (threadIdx.x < col_size) { 35 | block_boxes[threadIdx.x * 5 + 0] = 36 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; 37 | block_boxes[threadIdx.x * 5 + 1] = 38 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; 39 | block_boxes[threadIdx.x * 5 + 2] = 40 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; 41 | block_boxes[threadIdx.x * 5 + 3] = 42 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; 43 | block_boxes[threadIdx.x * 5 + 4] = 44 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; 45 | } 46 | __syncthreads(); 47 | 48 | if (threadIdx.x < row_size) { 49 | const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; 50 | const T* cur_box = dev_boxes + cur_box_idx * 5; 51 | int i = 0; 52 | unsigned long long t = 0; 53 | int start = 0; 54 | if (row_start == col_start) { 55 | start = threadIdx.x + 1; 56 | } 57 | for (i = start; i < col_size; i++) { 58 | // Instead of devIoU used by original horizontal nms, here 59 | // we use the single_box_iou_rotated function from box_iou_rotated_utils.h 60 | if (single_box_iou_rotated(cur_box, block_boxes + i * 5) > 61 | iou_threshold) { 62 | t |= 1ULL << i; 63 | } 64 | } 65 | const int col_blocks = at::cuda::ATenCeilDiv(n_boxes, threadsPerBlock); 66 | dev_mask[cur_box_idx * col_blocks + col_start] = t; 67 | } 68 | } 69 | 70 | 71 | at::Tensor nms_rotated_cuda( 72 | // input must be contiguous 73 | const at::Tensor& dets, 74 | const at::Tensor& scores, 75 | float iou_threshold) { 76 | // using scalar_t = float; 77 | AT_ASSERTM(dets.is_cuda(), "dets must be a CUDA tensor"); 78 | AT_ASSERTM(scores.is_cuda(), "scores must be a CUDA tensor"); 79 | at::cuda::CUDAGuard device_guard(dets.device()); 80 | 81 | auto order_t = std::get<1>(scores.sort(0, /* descending=*/true)); 82 | auto dets_sorted = dets.index_select(0, order_t); 83 | 84 | auto dets_num = dets.size(0); 85 | 86 | const int col_blocks = 87 | at::cuda::ATenCeilDiv(static_cast(dets_num), threadsPerBlock); 88 | 89 | at::Tensor mask = 90 | at::empty({dets_num * col_blocks}, dets.options().dtype(at::kLong)); 91 | 92 | dim3 blocks(col_blocks, col_blocks); 93 | dim3 threads(threadsPerBlock); 94 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 95 | 96 | AT_DISPATCH_FLOATING_TYPES( 97 | dets_sorted.scalar_type(), "nms_rotated_kernel_cuda", [&] { 98 | nms_rotated_cuda_kernel<<>>( 99 | dets_num, 100 | iou_threshold, 101 | dets_sorted.data_ptr(), 102 | (unsigned long long*)mask.data_ptr()); 103 | }); 104 | 105 | at::Tensor mask_cpu = mask.to(at::kCPU); 106 | unsigned long long* mask_host = 107 | (unsigned long long*)mask_cpu.data_ptr(); 108 | 109 | std::vector remv(col_blocks); 110 | memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks); 111 | 112 | at::Tensor keep = 113 | at::empty({dets_num}, dets.options().dtype(at::kLong).device(at::kCPU)); 114 | int64_t* keep_out = keep.data_ptr(); 115 | 116 | int num_to_keep = 0; 117 | for (int i = 0; i < dets_num; i++) { 118 | int nblock = i / threadsPerBlock; 119 | int inblock = i % threadsPerBlock; 120 | 121 | if (!(remv[nblock] & (1ULL << inblock))) { 122 | keep_out[num_to_keep++] = i; 123 | unsigned long long* p = mask_host + i * col_blocks; 124 | for (int j = nblock; j < col_blocks; j++) { 125 | remv[j] |= p[j]; 126 | } 127 | } 128 | } 129 | 130 | AT_CUDA_CHECK(cudaGetLastError()); 131 | return order_t.index( 132 | {keep.narrow(/*dim=*/0, /*start=*/0, /*length=*/num_to_keep) 133 | .to(order_t.device(), keep.scalar_type())}); 134 | } 135 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/src/nms_rotated_ext.cpp: -------------------------------------------------------------------------------- 1 | // Modified from 2 | // https://github.com/facebookresearch/detectron2/tree/master/detectron2/layers/csrc/nms_rotated 3 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved 4 | #include 5 | #include 6 | 7 | 8 | #ifdef WITH_CUDA 9 | at::Tensor nms_rotated_cuda( 10 | const at::Tensor& dets, 11 | const at::Tensor& scores, 12 | const float iou_threshold); 13 | 14 | at::Tensor poly_nms_cuda( 15 | const at::Tensor boxes, 16 | float nms_overlap_thresh); 17 | #endif 18 | 19 | at::Tensor nms_rotated_cpu( 20 | const at::Tensor& dets, 21 | const at::Tensor& scores, 22 | const float iou_threshold); 23 | 24 | 25 | inline at::Tensor nms_rotated( 26 | const at::Tensor& dets, 27 | const at::Tensor& scores, 28 | const float iou_threshold) { 29 | assert(dets.device().is_cuda() == scores.device().is_cuda()); 30 | if (dets.device().is_cuda()) { 31 | #ifdef WITH_CUDA 32 | return nms_rotated_cuda( 33 | dets.contiguous(), scores.contiguous(), iou_threshold); 34 | #else 35 | AT_ERROR("Not compiled with GPU support"); 36 | #endif 37 | } 38 | return nms_rotated_cpu(dets.contiguous(), scores.contiguous(), iou_threshold); 39 | } 40 | 41 | 42 | inline at::Tensor nms_poly( 43 | const at::Tensor& dets, 44 | const float iou_threshold) { 45 | if (dets.device().is_cuda()) { 46 | #ifdef WITH_CUDA 47 | if (dets.numel() == 0) 48 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 49 | return poly_nms_cuda(dets, iou_threshold); 50 | #else 51 | AT_ERROR("POLY_NMS is not compiled with GPU support"); 52 | #endif 53 | } 54 | AT_ERROR("POLY_NMS is not implemented on CPU"); 55 | } 56 | 57 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 58 | m.def("nms_rotated", &nms_rotated, "nms for rotated bboxes"); 59 | m.def("nms_poly", &nms_poly, "nms for poly bboxes"); 60 | } 61 | -------------------------------------------------------------------------------- /yolov5_obb/utils/nms_rotated/src/poly_nms_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | template 4 | at::Tensor poly_nms_cpu_kernel(const at::Tensor& dets, const float threshold) { 5 | 6 | -------------------------------------------------------------------------------- /yolov5_obb/xml2txt.py: -------------------------------------------------------------------------------- 1 | # *_* coding : UTF-8 *_* 2 | # 功能描述 ：把旋转框 cx,cy,w,h,angle，转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4,class,difficulty 3 | 4 | import os 5 | import xml.etree.ElementTree as ET 6 | import math 7 | 8 | label=['tomato'] 9 | def edit_xml(xml_file): 10 | """ 11 | 修改xml文件 12 | :param xml_file:xml文件的路径 13 | :return: 14 | """ 15 | print(xml_file) 16 | tree = ET.parse(xml_file) 17 | f=open(xml_file.replace('xml','txt').replace('anns','labelTxt'),'w') 18 | objs = tree.findall('object') 19 | for ix, obj in enumerate(objs): 20 | obj_type = obj.find('type') 21 | type = obj_type.text 22 | 23 | if type == 'bndbox': 24 | obj_bnd = obj.find('bndbox') 25 | obj_xmin = obj_bnd.find('xmin') 26 | obj_ymin = obj_bnd.find('ymin') 27 | obj_xmax = obj_bnd.find('xmax') 28 | obj_ymax = obj_bnd.find('ymax') 29 | xmin = float(obj_xmin.text) 30 | ymin = float(obj_ymin.text) 31 | xmax = float(obj_xmax.text) 32 | ymax = float(obj_ymax.text) 33 | obj_bnd.remove(obj_xmin) # 删除节点 34 | obj_bnd.remove(obj_ymin) 35 | obj_bnd.remove(obj_xmax) 36 | obj_bnd.remove(obj_ymax) 37 | x0 = xmin 38 | y0 = ymin 39 | x1 = xmax 40 | y1 = ymin 41 | x2 = xmin 42 | y2 = ymax 43 | x3 = xmax 44 | y3 = ymax 45 | elif type == 'robndbox': 46 | obj_bnd = obj.find('robndbox') 47 | obj_bnd.tag = 'bndbox' # 修改节点名 48 | obj_cx = obj_bnd.find('cx') 49 | obj_cy = obj_bnd.find('cy') 50 | obj_w = obj_bnd.find('w') 51 | obj_h = obj_bnd.find('h') 52 | obj_angle = obj_bnd.find('angle') 53 | cx = float(obj_cx.text) 54 | cy = float(obj_cy.text) 55 | w = float(obj_w.text) 56 | h = float(obj_h.text) 57 | angle = float(obj_angle.text) 58 | 59 | x0, y0 = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle) 60 | x1, y1 = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle) 61 | x2, y2 = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle) 62 | x3, y3 = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle) 63 | classes=int(obj.find('name').text) 64 | axis=list([str(x0),str(y0),str(x1), str(y1),str(x2), str(y2),str(x3), str(y3),label[classes],'0']) 65 | bb = " ".join(axis) 66 | f.writelines(bb) 67 | f.writelines("\n") 68 | f.close() 69 | # 转换成四点坐标 70 | def rotatePoint(xc, yc, xp, yp, theta): 71 | xoff = xp - xc; 72 | yoff = yp - yc; 73 | cosTheta = math.cos(theta) 74 | sinTheta = math.sin(theta) 75 | pResx = cosTheta * xoff + sinTheta * yoff 76 | pResy = - sinTheta * xoff + cosTheta * yoff 77 | return int(xc + pResx), int(yc + pResy) 78 | 79 | if __name__ == '__main__': 80 | for path in os.listdir('anns/'): 81 | edit_xml('anns/'+path) 82 | -------------------------------------------------------------------------------- /yolov5_obb/yolov5s.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/yolov5s.pt -------------------------------------------------------------------------------- /yolov5_obb/注意,根目录这个yolov5s.pt是预训练模型,不是最后旋转框的模型: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb/注意,根目录这个yolov5s.pt是预训练模型,不是最后旋转框的模型 -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.10) 2 | 3 | project(yolov5_obb) 4 | 5 | add_definitions(-std=c++11) 6 | add_definitions(-DAPI_EXPORTS) 7 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 8 | set(CMAKE_CXX_STANDARD 11) 9 | set(CMAKE_BUILD_TYPE Release) 10 | 11 | find_package(CUDA REQUIRED) 12 | 13 | if(WIN32) 14 | enable_language(CUDA) 15 | endif(WIN32) 16 | 17 | 18 | 19 | include_directories(${PROJECT_SOURCE_DIR}/include) 20 | find_package(OpenCV) 21 | include_directories(${OpenCV_INCLUDE_DIRS}) 22 | 23 | find_package(CUDA REQUIRED) 24 | message(STATUS " libraries: ${CUDA_LIBRARIES}") 25 | message(STATUS " include path: ${CUDA_INCLUDE_DIRS}") 26 | include_directories(${CUDA_INCLUDE_DIRS}) 27 | enable_language(CUDA) 28 | 29 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 30 | # cuda 31 | #include_directories(/usr/local/cuda/include) 32 | #link_directories(/usr/local/cuda/lib64) 33 | # tensorrt 34 | #include_directories(/usr/include/x86_64-linux-gnu/) 35 | #link_directories(/usr/lib/x86_64-linux-gnu/) 36 | 37 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -g -D_MWAITXINTRIN_H_INCLUDED")#-Wfatal-errors 38 | cuda_add_library(myplugins SHARED yololayer.cu) 39 | target_link_libraries(myplugins nvinfer cudart) 40 | 41 | 42 | cuda_add_executable(yolov5_gen yolov5_gen.cpp) 43 | 44 | target_link_libraries(yolov5_gen nvinfer) 45 | target_link_libraries(yolov5_gen cudart) 46 | target_link_libraries(yolov5_gen myplugins) 47 | target_link_libraries(yolov5_gen ${OpenCV_LIBS}) 48 | 49 | cuda_add_executable(yolov5_use yolov5_use.cpp) 50 | 51 | target_link_libraries(yolov5_use nvinfer) 52 | target_link_libraries(yolov5_use cudart) 53 | target_link_libraries(yolov5_use myplugins) 54 | target_link_libraries(yolov5_use ${OpenCV_LIBS}) 55 | 56 | 57 | if(UNIX) 58 | add_definitions(-O2 -pthread) 59 | endif(UNIX) 60 | 61 | 62 | -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/README.md: -------------------------------------------------------------------------------- 1 | # yolov5_for_oriented_object_detection 2 | 3 | The Pytorch implementation is [yolov5_obb](https://github.com/hukaixuan19970627/yolov5_obb ) 4 | 5 | actually ,This repo is based on [yolov5](https://github.com/ultralytics/yolov5/tree/v6.0/) and only used at Linux PC (partial dependency install hardly at windows PC) 6 | 7 | ## How to Run, yolov5s as example 8 | 9 | 1. generate .wts from pytorch with .pt 10 | 11 | ``` 12 | // clone code according to above #Different versions of yolov5 13 | train and gen best.pt at runs/train/exp[]/weights/ 14 | cp gen_wts.py {yolov5_obb} 15 | python gen_wts.py -w runs/train/exp[]/weights/best.pt -o yolov5s.wts 16 | // a file 'yolov5s.wts' will be generated. 17 | ``` 18 | 19 | 2. build **tensorr engine ** 20 | 21 | ``` 22 | // update CLASS_NUM in yololayer.h 23 | // caution:CLASS_NUM= your classes +180(180 for angle classes) 24 | mkdir build 25 | cd build 26 | cp {yolov5_obb}/yolov5s.wts ../yolov5s.wts 27 | cmake .. 28 | make 29 | sudo ./yolov5_gen -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file 30 | // For example yolov5s 31 | sudo ./yolov5_gen -s ../yolov5s.wts ../yolov5s.engine s 32 | ``` 33 | 34 | 3. use **tensorr engine** 35 | 36 | 37 | ``` 38 | sudo ./yolov5_use ../../yolov5s.engine ../../images/mytest.png 39 | ``` 40 | 41 | ![image-20220402142552499](C:\Users\Administrator\AppData\Roaming\Typora\typora-user-images\image-20220402142552499.png) 42 | 43 | 44 | 45 | -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef TRTX_CUDA_UTILS_H_ 2 | #define TRTX_CUDA_UTILS_H_ 3 | 4 | #include 5 | 6 | #ifndef CUDA_CHECK 7 | #define CUDA_CHECK(callstr)\ 8 | {\ 9 | cudaError_t error_code = callstr;\ 10 | if (error_code != cudaSuccess) {\ 11 | std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\ 12 | assert(0);\ 13 | }\ 14 | } 15 | #endif // CUDA_CHECK 16 | 17 | #endif // TRTX_CUDA_UTILS_H_ 18 | 19 | -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/gen_wts.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import argparse 3 | import os 4 | import struct 5 | import torch 6 | from utils.torch_utils import select_device 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser(description='Convert .pt file to .wts') 11 | parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)') 12 | parser.add_argument('-o', '--output', help='Output (.wts) file path (optional)') 13 | args = parser.parse_args() 14 | if not os.path.isfile(args.weights): 15 | raise SystemExit('Invalid input file') 16 | if not args.output: 17 | args.output = os.path.splitext(args.weights)[0] + '.wts' 18 | elif os.path.isdir(args.output): 19 | args.output = os.path.join( 20 | args.output, 21 | os.path.splitext(os.path.basename(args.weights))[0] + '.wts') 22 | return args.weights, args.output 23 | 24 | 25 | pt_file, wts_file = parse_args() 26 | 27 | # Initialize 28 | device = select_device('cpu') 29 | # Load model 30 | model = torch.load(pt_file, map_location=device) # load to FP32 31 | model = model['ema' if model.get('ema') else 'model'].float() 32 | 33 | # update anchor_grid info 34 | anchor_grid = model.model[-1].anchors * model.model[-1].stride[...,None,None] 35 | # model.model[-1].anchor_grid = anchor_grid 36 | delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer 37 | model.model[-1].register_buffer("anchor_grid",anchor_grid) #The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight. 38 | 39 | model.to(device).eval() 40 | 41 | with open(wts_file, 'w') as f: 42 | f.write('{}\n'.format(len(model.state_dict().keys()))) 43 | for k, v in model.state_dict().items(): 44 | vr = v.reshape(-1).cpu().numpy() 45 | f.write('{} {} '.format(k, len(vr))) 46 | for vv in vr: 47 | f.write(' ') 48 | f.write(struct.pack('>f' ,float(vv)).hex()) 49 | f.write('\n') 50 | -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/images/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb_tensorrt_cpp/images/test.jpg -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/macros.h: -------------------------------------------------------------------------------- 1 | #ifndef __MACROS_H 2 | #define __MACROS_H 3 | 4 | #ifdef API_EXPORTS 5 | #if defined(_MSC_VER) 6 | #define API __declspec(dllexport) 7 | #else 8 | #define API __attribute__((visibility("default"))) 9 | #endif 10 | #else 11 | 12 | #if defined(_MSC_VER) 13 | #define API __declspec(dllimport) 14 | #else 15 | #define API 16 | #endif 17 | #endif // API_EXPORTS 18 | 19 | #if NV_TENSORRT_MAJOR >= 8 20 | #define TRT_NOEXCEPT noexcept 21 | #define TRT_CONST_ENQUEUE const 22 | #else 23 | #define TRT_NOEXCEPT 24 | #define TRT_CONST_ENQUEUE 25 | #endif 26 | 27 | #endif // __MACROS_H 28 | -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb_tensorrt_cpp/result.jpg -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/yololayer.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb_tensorrt_cpp/yololayer.cu -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/yololayer.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb_tensorrt_cpp/yololayer.h -------------------------------------------------------------------------------- /yolov5_obb_tensorrt_cpp/yolov5_use.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fish-kong/Yolov5-obb-Tensorrt-Infer/dffc8b4d330e2d51875898e8b96582b19ef4a8c2/yolov5_obb_tensorrt_cpp/yolov5_use.cpp --------------------------------------------------------------------------------