├── Arial.ttf ├── CONTRIBUTING.md ├── 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 ├── __init__.py ├── __pycache__ │ ├── dota_poly2rbox.cpython-37.pyc │ └── dota_utils.cpython-37.pyc ├── _polyiou.cpython-37m-x86_64-linux-gnu.so ├── build │ ├── lib.linux-x86_64-cpython-37 │ │ └── _polyiou.cpython-37m-x86_64-linux-gnu.so │ └── temp.linux-x86_64-cpython-37 │ │ ├── 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 │ ├── 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 │ └── 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 ├── Dockerfile ├── LICENSE ├── README.md ├── YoloObbTrack ├── Arial.ttf ├── Tracker.py ├── __pycache__ │ ├── AIDetector.cpython-37.pyc │ ├── BaseDetector.cpython-37.pyc │ ├── Tracker.cpython-37.pyc │ └── Tracker.cpython-38.pyc ├── byte_tracker │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── byte_tracker.cpython-37.pyc │ │ ├── byte_tracker.cpython-38.pyc │ │ └── strong_sort.cpython-37.pyc │ ├── byte_tracker.py │ └── sort │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── detection.cpython-37.pyc │ │ ├── iou_matching.cpython-37.pyc │ │ ├── iou_matching.cpython-38.pyc │ │ ├── kalman_filter.cpython-37.pyc │ │ ├── kalman_filter_eight.cpython-37.pyc │ │ ├── kalman_filter_rbox.cpython-37.pyc │ │ ├── kalman_filter_rbox.cpython-38.pyc │ │ ├── linear_assignment.cpython-37.pyc │ │ ├── linear_assignment.cpython-38.pyc │ │ ├── nn_matching.cpython-37.pyc │ │ ├── track.cpython-37.pyc │ │ ├── track_eight.cpython-37.pyc │ │ ├── track_rbox.cpython-37.pyc │ │ ├── track_rbox.cpython-38.pyc │ │ ├── tracker.cpython-37.pyc │ │ └── tracker.cpython-38.pyc │ │ ├── iou_matching.py │ │ ├── kalman_filter_rbox.py │ │ ├── linear_assignment.py │ │ ├── track_rbox.py │ │ └── tracker.py └── strong_sort │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── strong_sort.cpython-37.pyc │ └── strong_sort.cpython-38.pyc │ ├── configs │ └── strong_sort.yaml │ ├── deep │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── feature_extractor.cpython-37.pyc │ │ ├── feature_extractor.cpython-38.pyc │ │ ├── model.cpython-37.pyc │ │ └── model.cpython-38.pyc │ ├── checkpoint │ │ ├── .gitkeep │ │ └── osnet_x0_25_msmt17.pth │ ├── feature_extractor.py │ └── model.py │ ├── sort │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── detection.cpython-37.pyc │ │ ├── detection.cpython-38.pyc │ │ ├── iou_matching.cpython-37.pyc │ │ ├── iou_matching.cpython-38.pyc │ │ ├── kalman_filter.cpython-37.pyc │ │ ├── kalman_filter_eight.cpython-37.pyc │ │ ├── kalman_filter_rbox.cpython-37.pyc │ │ ├── kalman_filter_rbox.cpython-38.pyc │ │ ├── linear_assignment.cpython-37.pyc │ │ ├── linear_assignment.cpython-38.pyc │ │ ├── nn_matching.cpython-37.pyc │ │ ├── nn_matching.cpython-38.pyc │ │ ├── track.cpython-37.pyc │ │ ├── track_eight.cpython-37.pyc │ │ ├── track_rbox.cpython-37.pyc │ │ ├── track_rbox.cpython-38.pyc │ │ ├── tracker.cpython-37.pyc │ │ └── tracker.cpython-38.pyc │ ├── detection.py │ ├── iou_matching.py │ ├── kalman_filter_rbox.py │ ├── linear_assignment.py │ ├── nn_matching.py │ ├── preprocessing.py │ ├── track_rbox.py │ └── tracker.py │ ├── strong_sort.py │ └── utils │ ├── __init__.py │ ├── asserts.py │ ├── draw.py │ ├── evaluation.py │ ├── io.py │ ├── json_logger.py │ ├── log.py │ ├── parser.py │ └── tools.py ├── __pycache__ ├── val.cpython-37.pyc └── val.cpython-38.pyc ├── data ├── dotav1_poly.yaml ├── hyps │ └── obb │ │ ├── hyp.finetune_DroneVehicle.yaml │ │ ├── hyp.finetune_dota.yaml │ │ ├── hyp.finetune_dota_CloseAug.yaml │ │ ├── hyp.finetune_dota_kpt.yaml │ │ └── hyp.paper.yaml ├── scripts │ └── download_weights.sh ├── yolov8obb_demo.yaml ├── yolov8obb_demo_split.yaml └── yolov8obb_kpt_demo.yaml ├── datasets └── uav │ ├── images │ └── train │ │ ├── DJI_0007_006750.jpg │ │ ├── DJI_0007_006840.jpg │ │ ├── DJI_0007_006900.jpg │ │ ├── DJI_0007_006990.jpg │ │ ├── DJI_0007_007020.jpg │ │ ├── DJI_0007_007140.jpg │ │ ├── DJI_0007_007200.jpg │ │ ├── DJI_0007_007320.jpg │ │ ├── DJI_0007_007410.jpg │ │ ├── DJI_0007_007470.jpg │ │ ├── DJI_0007_007560.jpg │ │ ├── DJI_0007_009090.jpg │ │ ├── DJI_0007_009960.jpg │ │ ├── DJI_0007_010800.jpg │ │ ├── DJI_0007_010890.jpg │ │ ├── DJI_0007_010920.jpg │ │ ├── DJI_0007_011010.jpg │ │ ├── DJI_0007_011070.jpg │ │ ├── DJI_0007_011160.jpg │ │ ├── DJI_0007_011190.jpg │ │ ├── DJI_0007_011280.jpg │ │ ├── DJI_0007_011340.jpg │ │ ├── DJI_0007_011430.jpg │ │ ├── DJI_0007_011460.jpg │ │ ├── DJI_0007_011550.jpg │ │ ├── DJI_0007_011610.jpg │ │ ├── DJI_0012_000000.jpg │ │ ├── DJI_0012_000270.jpg │ │ ├── DJI_0012_000540.jpg │ │ ├── DJI_0021_000870.jpg │ │ ├── DJI_0021_001140.jpg │ │ ├── DJI_0021_001410.jpg │ │ ├── DJI_0021_001680.jpg │ │ ├── DJI_0021_001950.jpg │ │ ├── DJI_0021_002490.jpg │ │ ├── DJI_0021_003030.jpg │ │ ├── DJI_0021_003300.jpg │ │ ├── DJI_0021_003570.jpg │ │ ├── DJI_0021_003840.jpg │ │ └── DJI_0039_004200.jpg │ ├── labelTxt │ └── train │ │ ├── DJI_0007_006750.txt │ │ ├── DJI_0007_006840.txt │ │ ├── DJI_0007_006900.txt │ │ ├── DJI_0007_006990.txt │ │ ├── DJI_0007_007020.txt │ │ ├── DJI_0007_007140.txt │ │ ├── DJI_0007_007200.txt │ │ ├── DJI_0007_007320.txt │ │ ├── DJI_0007_007410.txt │ │ ├── DJI_0007_007470.txt │ │ ├── DJI_0007_007560.txt │ │ ├── DJI_0007_009090.txt │ │ ├── DJI_0007_009960.txt │ │ ├── DJI_0007_010800.txt │ │ ├── DJI_0007_010890.txt │ │ ├── DJI_0007_010920.txt │ │ ├── DJI_0007_011010.txt │ │ ├── DJI_0007_011070.txt │ │ ├── DJI_0007_011160.txt │ │ ├── DJI_0007_011190.txt │ │ ├── DJI_0007_011280.txt │ │ ├── DJI_0007_011340.txt │ │ ├── DJI_0007_011430.txt │ │ ├── DJI_0007_011460.txt │ │ ├── DJI_0007_011550.txt │ │ ├── DJI_0007_011610.txt │ │ ├── DJI_0012_000000.txt │ │ ├── DJI_0012_000270.txt │ │ ├── DJI_0012_000540.txt │ │ ├── DJI_0021_000870.txt │ │ ├── DJI_0021_001140.txt │ │ ├── DJI_0021_001410.txt │ │ ├── DJI_0021_001680.txt │ │ ├── DJI_0021_001950.txt │ │ ├── DJI_0021_002490.txt │ │ ├── DJI_0021_003030.txt │ │ ├── DJI_0021_003300.txt │ │ ├── DJI_0021_003570.txt │ │ ├── DJI_0021_003840.txt │ │ └── DJI_0039_004200.txt │ ├── train.cache │ ├── train.txt │ ├── val.cache │ └── val.txt ├── detect.py ├── docs ├── GetStart.md ├── install.md └── obb_kpt_result.jpg ├── eval_rotate_PR_V8.py ├── export.py ├── models ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── common.cpython-37.pyc │ ├── common.cpython-38.pyc │ ├── common_prune.cpython-37.pyc │ ├── common_prune.cpython-38.pyc │ ├── experimental.cpython-37.pyc │ ├── experimental.cpython-38.pyc │ ├── yolo.cpython-37.pyc │ ├── yolo.cpython-38.pyc │ ├── yolo_prune.cpython-37.pyc │ └── yolo_prune.cpython-38.pyc ├── common.py ├── common_prune.py ├── experimental.py ├── hub │ ├── yolov8n_C2FTS.yaml │ ├── yolov8n_C3TS.yaml │ ├── yolov8n_RFCAConv.yaml │ ├── yolov8n_c2fTS_cbam.yaml │ ├── yolov8n_ca.yaml │ └── yolov8n_cbam.yaml ├── yaml │ ├── yolov8l.yaml │ ├── yolov8n.yaml │ ├── yolov8n_kpt.yaml │ ├── yolov8nx6.yaml │ ├── yolov8s.yaml │ └── yolov8x.yaml ├── yolo.py └── yolo_prune.py ├── prune.py ├── prune_finetune.py ├── requirements.txt ├── setup.cfg ├── sh └── ddp_train.sh ├── tools ├── TestJson2VocClassTxt.py ├── Xml2Txt.py ├── hubconf.py ├── json2obb_kpt_labels.py └── mk_train.py ├── track_predict.py ├── train.py ├── utils ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── activations.cpython-37.pyc │ ├── augmentations.cpython-37.pyc │ ├── augmentations.cpython-38.pyc │ ├── autoanchor.cpython-37.pyc │ ├── autoanchor.cpython-38.pyc │ ├── autobatch.cpython-37.pyc │ ├── autobatch.cpython-38.pyc │ ├── callbacks.cpython-37.pyc │ ├── callbacks.cpython-38.pyc │ ├── datasets.cpython-37.pyc │ ├── datasets.cpython-38.pyc │ ├── downloads.cpython-37.pyc │ ├── downloads.cpython-38.pyc │ ├── general.cpython-37.pyc │ ├── general.cpython-38.pyc │ ├── lion_pytorch.cpython-37.pyc │ ├── lion_pytorch.cpython-38.pyc │ ├── loss.cpython-37.pyc │ ├── loss.cpython-38.pyc │ ├── metrics.cpython-37.pyc │ ├── metrics.cpython-38.pyc │ ├── plots.cpython-37.pyc │ ├── plots.cpython-38.pyc │ ├── prune_utils.cpython-37.pyc │ ├── prune_utils.cpython-38.pyc │ ├── rboxs_utils.cpython-37.pyc │ ├── rboxs_utils.cpython-38.pyc │ ├── tal.cpython-37.pyc │ ├── tal.cpython-38.pyc │ ├── torch_utils.cpython-37.pyc │ └── torch_utils.cpython-38.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 ├── lion_pytorch.py ├── loggers │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── __init__.cpython-38.pyc │ └── wandb │ │ ├── README.md │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── wandb_utils.cpython-37.pyc │ │ ├── wandb_utils.cpython-38.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 │ │ ├── __init__.cpython-38.pyc │ │ ├── nms_rotated_wrapper.cpython-37.pyc │ │ └── nms_rotated_wrapper.cpython-38.pyc │ ├── build │ │ ├── lib.linux-x86_64-cpython-37 │ │ │ └── nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so │ │ ├── temp.linux-x86_64-cpython-37 │ │ │ └── src │ │ │ │ ├── nms_rotated_cpu.o │ │ │ │ ├── nms_rotated_cuda.o │ │ │ │ ├── nms_rotated_ext.o │ │ │ │ └── poly_nms_cuda.o │ │ └── temp.linux-x86_64-cpython-38 │ │ │ └── src │ │ │ ├── nms_rotated_cpu.o │ │ │ ├── nms_rotated_cuda.o │ │ │ └── nms_rotated_ext.o │ ├── nms_rotated.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ ├── not-zip-safe │ │ └── top_level.txt │ ├── nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so │ ├── 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 ├── prune_utils.py ├── rboxs_utils.py ├── tal.py └── torch_utils.py ├── val.py └── weights └── yolov8n.pt /Arial.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/Arial.ttf -------------------------------------------------------------------------------- /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) -------------------------------------------------------------------------------- /DOTA_devkit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/__init__.py -------------------------------------------------------------------------------- /DOTA_devkit/__pycache__/dota_poly2rbox.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/__pycache__/dota_poly2rbox.cpython-37.pyc -------------------------------------------------------------------------------- /DOTA_devkit/__pycache__/dota_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/__pycache__/dota_utils.cpython-37.pyc -------------------------------------------------------------------------------- /DOTA_devkit/_polyiou.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/_polyiou.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /DOTA_devkit/build/lib.linux-x86_64-cpython-37/_polyiou.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/build/lib.linux-x86_64-cpython-37/_polyiou.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /DOTA_devkit/build/temp.linux-x86_64-cpython-37/polyiou.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/build/temp.linux-x86_64-cpython-37/polyiou.o -------------------------------------------------------------------------------- /DOTA_devkit/build/temp.linux-x86_64-cpython-37/polyiou_wrap.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/build/temp.linux-x86_64-cpython-37/polyiou_wrap.o -------------------------------------------------------------------------------- /DOTA_devkit/poly_nms_gpu/Makefile: -------------------------------------------------------------------------------- 1 | all: 2 | python setup.py build_ext --inplace 3 | rm -rf build 4 | -------------------------------------------------------------------------------- /DOTA_devkit/poly_nms_gpu/__init__.py: -------------------------------------------------------------------------------- 1 | from .poly_overlaps import poly_overlaps 2 | 3 | __all__ = ['poly_overlaps', 'poly_nms'] 4 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /DOTA_devkit/poly_nms_gpu/poly_nms_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/DOTA_devkit/poly_nms_gpu/poly_nms_test.py -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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() -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch 4 | FROM nvcr.io/nvidia/pytorch:21.10-py3 5 | 6 | # Install linux packages 7 | RUN apt update && apt install -y zip htop screen libgl1-mesa-glx 8 | 9 | # Install python dependencies 10 | COPY requirements.txt . 11 | RUN python -m pip install --upgrade pip 12 | RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof 13 | RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook wandb>=0.12.2 14 | RUN pip install --no-cache -U torch torchvision numpy Pillow 15 | # RUN pip install --no-cache torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html 16 | 17 | # Create working directory 18 | RUN mkdir -p /usr/src/app 19 | WORKDIR /usr/src/app 20 | 21 | # Copy contents 22 | COPY . /usr/src/app 23 | 24 | # Downloads to user config dir 25 | ADD https://ultralytics.com/assets/Arial.ttf /root/.config/Ultralytics/ 26 | 27 | # Set environment variables 28 | # ENV HOME=/usr/src/app 29 | 30 | 31 | # Usage Examples ------------------------------------------------------------------------------------------------------- 32 | 33 | # Build and Push 34 | # t=ultralytics/yolov5:latest && sudo docker build -t $t . && sudo docker push $t 35 | 36 | # Pull and Run 37 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t 38 | 39 | # Pull and Run with local directory access 40 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t 41 | 42 | # Kill all 43 | # sudo docker kill $(sudo docker ps -q) 44 | 45 | # Kill all image-based 46 | # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest) 47 | 48 | # Bash into running container 49 | # sudo docker exec -it 5a9b5863d93d bash 50 | 51 | # Bash into stopped container 52 | # id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash 53 | 54 | # Clean up 55 | # docker system prune -a --volumes 56 | 57 | # Update Ubuntu drivers 58 | # https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/ 59 | 60 | # DDP test 61 | # python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3 62 | 63 | # GCP VM from Image 64 | # docker.io/ultralytics/yolov5:latest 65 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Yolov8 for Oriented Object Detection 2 | ![图片](./docs/obb_kpt_result.jpg) 3 | 4 | 5 | # Installation 6 | Please refer to [install.md](./docs/install.md) for installation and dataset preparation. 7 | 8 | # Getting Started 9 | This repo is based on [yolov5_obb](https://github.com/hukaixuan19970627/yolov5_obb/tree/master). [yolov5_obb_prune_tracking](https://github.com/yzqxy/yolov5_obb_prune_tracking) 10 | 11 | And this repo has been rebuilt, Please see [GetStart.md](./docs/GetStart.md) for the Oriented Detection latest basic usage. 12 | 13 | # Acknowledgements 14 | I have used utility functions from other wonderful open-source projects. Espeicially thank the authors of: 15 | 16 | * [ultralytics/yolov5](https://github.com/ultralytics/yolov5). 17 | * [Thinklab-SJTU/CSL_RetinaNet_Tensorflow](https://github.com/Thinklab-SJTU/CSL_RetinaNet_Tensorflow). 18 | * [jbwang1997/OBBDetection](https://github.com/jbwang1997/OBBDetection) 19 | * [CAPTAIN-WHU/DOTA_devkit](https://github.com/CAPTAIN-WHU/DOTA_devkit) 20 | * [yolov5_obb](https://github.com/hukaixuan19970627/yolov5_obb/tree/master) 21 | * [yolov5_obb_prune_tracking](https://github.com/yzqxy/yolov5_obb_prune_tracking/tree/master) 22 | ## More detailed explanation 23 | 想要了解相关实现的细节和原理可以看我的CSDN文章: 24 | * [Yolov8_obb基于anchor_free的旋转框目标检测，剪枝，跟踪](https://blog.csdn.net/qq_39128381/article/details/131962684?spm=1001.2014.3001.5501). 25 | 26 | ## 有问题反馈 27 | 在使用中有任何问题，建议先按照[install.md](./docs/install.md)检查环境依赖项，再按照[GetStart.md](./docs/GetStart.md)检查使用流程是否正确，善用搜索引擎和github中的issue搜索框，可以极大程度上节省你的时间。 28 | 29 | * 代码问题提issues,其他问题请CSDN上联系 30 | 31 | -------------------------------------------------------------------------------- /YoloObbTrack/Arial.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/Arial.ttf -------------------------------------------------------------------------------- /YoloObbTrack/Tracker.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Author: yzqxy 3 | ''' 4 | from .strong_sort.strong_sort import StrongSORT 5 | from .byte_tracker.byte_tracker import Byte_tracker 6 | import torch 7 | import numpy as np 8 | from PIL import ImageDraw, ImageFont 9 | from utils.rboxs_utils import rbox2poly,poly2hbb 10 | palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) 11 | deepsort = StrongSORT() 12 | byte_tracker = Byte_tracker() 13 | 14 | 15 | 16 | def scale_polys(img1_shape, polys, img0_shape, ratio_pad=None): 17 | if ratio_pad is None: # calculate from img0_shape 18 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = resized / raw 19 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding 20 | else: 21 | gain = ratio_pad[0][0] # h_ratios 22 | pad = ratio_pad[1] # wh_paddings 23 | 24 | polys[:, [0, 2, 4, 6]] -= pad[0] # x padding 25 | polys[:, [1, 3, 5, 7]] -= pad[1] # y padding 26 | polys[:, :8] /= gain # Rescale poly shape to img0_shape 27 | #clip_polys(polys, img0_shape) 28 | return polys 29 | 30 | def scale_polys_single_img(img1_shape, polys, img0_shape, ratio_pad=None): 31 | if ratio_pad is None: # calculate from img0_shape 32 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = resized / raw 33 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding 34 | else: 35 | gain = ratio_pad[0][0] # h_ratios 36 | pad = ratio_pad[1] # wh_paddings 37 | 38 | polys[ [0, 2, 4, 6]] -= pad[0] # x padding 39 | polys[ [1, 3, 5, 7]] -= pad[1] # y padding 40 | polys[ :8] /= gain # Rescale poly shape to img0_shape 41 | #clip_polys(polys, img0_shape) 42 | return polys 43 | 44 | def compute_color_for_labels(label): 45 | """ 46 | Simple function that adds fixed color depending on the class 47 | """ 48 | color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette] 49 | return tuple(color) 50 | 51 | 52 | def draw_boxes(image, rbox, font,img,ori_img,track_type): 53 | 54 | for rrbox, cls_name, track_id in rbox: 55 | color = compute_color_for_labels(track_id) 56 | label = '{}{:d}'.format(cls_name, track_id) 57 | draw = ImageDraw.Draw(image) 58 | label = label.encode('utf-8') 59 | rrbox=rbox2poly(rrbox) 60 | rrbox=scale_polys_single_img(img,rrbox,ori_img) 61 | 62 | text_origin = np.array([rrbox[0], rrbox[1]], np.int32) 63 | draw.polygon(xy=list(rrbox), outline=color, width=2) 64 | draw.text(text_origin, str(label,'UTF-8'), fill=color, font=font) 65 | del draw 66 | return image 67 | 68 | 69 | 70 | 71 | def update_tracker( image,pred,PIL_image,img,ori_img,track_type='Byte_tracker'): 72 | 73 | top_label = np.array(pred[0][:, 6].cpu(), dtype = 'int32') 74 | rbox = pred[0][:, :5] 75 | confss =pred[0][:, 5] 76 | 77 | if track_type=='Byte_tracker': 78 | outputs = byte_tracker.update( rbox, confss, top_label, image) 79 | else: 80 | box=poly2hbb(rbox2poly(rbox)) 81 | outputs = deepsort.update( box.cpu(),rbox.cpu(), confss, top_label, image) 82 | 83 | print('outputs',outputs) 84 | rboxes2draw = [] 85 | current_ids = [] 86 | for value in list(outputs): 87 | rbox, track_id, cls_, _= value 88 | 89 | rboxes2draw.append((rbox, cls_, track_id)) 90 | current_ids.append(track_id) 91 | 92 | 93 | font = ImageFont.truetype(font='YoloObbTrack/Arial.ttf', 94 | size=np.floor(3e-2 * PIL_image.size[1] + 0.5).astype('int32')) 95 | image = draw_boxes(PIL_image, rboxes2draw, font,img,ori_img,track_type) 96 | 97 | 98 | return image, rboxes2draw 99 | -------------------------------------------------------------------------------- /YoloObbTrack/__pycache__/AIDetector.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/__pycache__/AIDetector.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/__pycache__/BaseDetector.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/__pycache__/BaseDetector.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/__pycache__/Tracker.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/__pycache__/Tracker.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/__pycache__/Tracker.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/__pycache__/Tracker.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/__init__.py: -------------------------------------------------------------------------------- 1 | from .byte_tracker import Byte_tracker 2 | 3 | 4 | __all__ = ['Byte_tracker', 'build_tracker'] 5 | 6 | 7 | def build_tracker(cfg, use_cuda): 8 | return Byte_tracker(cfg.STRONGSORT.REID_CKPT, 9 | max_dist=cfg.STRONGSORT.MAX_DIST, min_confidence=cfg.STRONGSORT.MIN_CONFIDENCE, 10 | nms_max_overlap=cfg.STRONGSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE, 11 | max_age=cfg.STRONGSORT.MAX_AGE, n_init=cfg.STRONGSORT.N_INIT, nn_budget=cfg.STRONGSORT.NN_BUDGET, use_cuda=use_cuda) 12 | -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/__pycache__/byte_tracker.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/__pycache__/byte_tracker.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/__pycache__/byte_tracker.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/__pycache__/byte_tracker.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/__pycache__/strong_sort.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/__pycache__/strong_sort.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/byte_tracker.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from .sort.tracker import Tracker 4 | 5 | __all__ = ['Byte_tracker'] 6 | 7 | 8 | class Byte_tracker(object): 9 | def __init__(self, 10 | max_iou_distance=0.5, 11 | max_age=70, 12 | n_init=3, 13 | ): 14 | 15 | self.tracker = Tracker( max_iou_distance=max_iou_distance, 16 | max_age=max_age, n_init=n_init) 17 | 18 | def update(self, rbox, confidences, classes, ori_img): 19 | self.height, self.width = ori_img.shape[:2] 20 | # 根据检测框的分数将框分为高质量框det和低质量框det_second,阈值可自行设定 21 | track_high_thresh=0.5 22 | track_low_thresh=0.1 23 | det=[] 24 | det_second=[] 25 | 26 | for i, conf in enumerate(confidences): 27 | if conf >= track_high_thresh: 28 | det.append([rbox[i],confidences[i],classes[i]]) 29 | elif conf > track_low_thresh : 30 | det_second.append([rbox[i],confidences[i],classes[i]]) 31 | 32 | # update tracker 33 | self.tracker.predict() 34 | self.tracker.update(det, det_second) 35 | 36 | # output bbox identities 37 | outputs = [] 38 | for track in self.tracker.tracks: 39 | #该track处于未证实状态或者是time_since_update大于3帧则跳出本次循环 40 | if not track.is_confirmed() or track.time_since_update > 3: 41 | continue 42 | rbox = track.mean[:5] 43 | track_id = track.track_id 44 | class_id = track.class_id 45 | conf = track.conf 46 | outputs.append([rbox, track_id, class_id, np.array(conf.cpu())]) 47 | if len(outputs) > 0: 48 | outputs = np.stack(outputs, axis=0) 49 | return outputs 50 | 51 | 52 | -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__init__.py -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/detection.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/detection.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/iou_matching.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/iou_matching.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/iou_matching.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/iou_matching.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter_eight.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter_eight.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter_rbox.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter_rbox.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter_rbox.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/kalman_filter_rbox.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/linear_assignment.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/linear_assignment.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/linear_assignment.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/linear_assignment.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/nn_matching.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/nn_matching.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/track.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/track.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/track_eight.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/track_eight.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/track_rbox.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/track_rbox.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/track_rbox.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/track_rbox.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/tracker.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/tracker.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/__pycache__/tracker.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/byte_tracker/sort/__pycache__/tracker.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/byte_tracker/sort/linear_assignment.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | from __future__ import absolute_import 3 | import numpy as np 4 | from scipy.optimize import linear_sum_assignment 5 | 6 | 7 | 8 | INFTY_COST = 1e+5 9 | 10 | 11 | def min_cost_matching( 12 | distance_metric, max_distance, tracks, detections, track_indices=None, 13 | detection_indices=None): 14 | """Solve linear assignment problem. 15 | Parameters 16 | ---------- 17 | distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray 18 | The distance metric is given a list of tracks and detections as well as 19 | a list of N track indices and M detection indices. The metric should 20 | return the NxM dimensional cost matrix, where element (i, j) is the 21 | association cost between the i-th track in the given track indices and 22 | the j-th detection in the given detection_indices. 23 | max_distance : float 24 | Gating threshold. Associations with cost larger than this value are 25 | disregarded. 26 | tracks : List[track.Track] 27 | A list of predicted tracks at the current time step. 28 | detections : List[detection.Detection] 29 | A list of detections at the current time step. 30 | track_indices : List[int] 31 | List of track indices that maps rows in `cost_matrix` to tracks in 32 | `tracks` (see description above). 33 | detection_indices : List[int] 34 | List of detection indices that maps columns in `cost_matrix` to 35 | detections in `detections` (see description above). 36 | Returns 37 | ------- 38 | (List[(int, int)], List[int], List[int]) 39 | Returns a tuple with the following three entries: 40 | * A list of matched track and detection indices. 41 | * A list of unmatched track indices. 42 | * A list of unmatched detection indices. 43 | """ 44 | if track_indices is None: 45 | track_indices = np.arange(len(tracks)) 46 | if detection_indices is None: 47 | detection_indices = np.arange(len(detections)) 48 | 49 | # if len(detection_indices) == 0 or len(track_indices) == 0: 50 | # return [], track_indices, detection_indices # Nothing to match. 51 | 52 | cost_matrix = distance_metric( 53 | tracks, detections, track_indices, detection_indices) 54 | 55 | cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5 56 | 57 | row_indices, col_indices = linear_sum_assignment(cost_matrix) 58 | 59 | matches, unmatched_tracks, unmatched_detections = [], [], [] 60 | for col, detection_idx in enumerate(detection_indices): 61 | if col not in col_indices: 62 | unmatched_detections.append(detection_idx) 63 | for row, track_idx in enumerate(track_indices): 64 | if row not in row_indices: 65 | unmatched_tracks.append(track_idx) 66 | for row, col in zip(row_indices, col_indices): 67 | track_idx = track_indices[row] 68 | detection_idx = detection_indices[col] 69 | if cost_matrix[row, col] > max_distance: 70 | unmatched_tracks.append(track_idx) 71 | unmatched_detections.append(detection_idx) 72 | else: 73 | matches.append((track_idx, detection_idx)) 74 | return matches, unmatched_tracks, unmatched_detections 75 | 76 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/__init__.py: -------------------------------------------------------------------------------- 1 | from .strong_sort import StrongSORT 2 | 3 | 4 | __all__ = ['StrongSORT', 'build_tracker'] 5 | 6 | 7 | def build_tracker(cfg, use_cuda): 8 | return StrongSORT(cfg.STRONGSORT.REID_CKPT, 9 | max_dist=cfg.STRONGSORT.MAX_DIST, min_confidence=cfg.STRONGSORT.MIN_CONFIDENCE, 10 | nms_max_overlap=cfg.STRONGSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE, 11 | max_age=cfg.STRONGSORT.MAX_AGE, n_init=cfg.STRONGSORT.N_INIT, nn_budget=cfg.STRONGSORT.NN_BUDGET, use_cuda=use_cuda) 12 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/__pycache__/strong_sort.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/__pycache__/strong_sort.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/__pycache__/strong_sort.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/__pycache__/strong_sort.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/configs/strong_sort.yaml: -------------------------------------------------------------------------------- 1 | STRONGSORT: 2 | ECC: True # activate camera motion compensation 3 | MC_LAMBDA: 0.995 # matching with both appearance (1 - MC_LAMBDA) and motion cost 4 | EMA_ALPHA: 0.9 # updates appearance state in an exponential moving average manner 5 | MAX_DIST: 0.2 # The matching threshold. Samples with larger distance are considered an invalid match 6 | MAX_IOU_DISTANCE: 0.7 # Gating threshold. Associations with cost larger than this value are disregarded. 7 | MAX_AGE: 30 # Maximum number of missed misses before a track is deleted 8 | N_INIT: 3 # Number of frames that a track remains in initialization phase 9 | NN_BUDGET: 100 # Maximum size of the appearance descriptors gallery 10 | 11 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__init__.py -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__pycache__/feature_extractor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__pycache__/feature_extractor.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__pycache__/feature_extractor.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__pycache__/feature_extractor.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__pycache__/model.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__pycache__/model.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/__pycache__/model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/__pycache__/model.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/checkpoint/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/checkpoint/.gitkeep -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/deep/checkpoint/osnet_x0_25_msmt17.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/deep/checkpoint/osnet_x0_25_msmt17.pth -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__init__.py -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/detection.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/detection.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/detection.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/detection.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/iou_matching.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/iou_matching.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/iou_matching.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/iou_matching.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter_eight.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter_eight.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter_rbox.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter_rbox.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter_rbox.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/kalman_filter_rbox.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/linear_assignment.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/linear_assignment.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/linear_assignment.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/linear_assignment.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/nn_matching.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/nn_matching.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/nn_matching.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/nn_matching.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/track.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/track.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/track_eight.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/track_eight.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/track_rbox.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/track_rbox.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/track_rbox.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/track_rbox.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/tracker.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/tracker.cpython-37.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/__pycache__/tracker.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/sort/__pycache__/tracker.cpython-38.pyc -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/detection.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | import numpy as np 3 | 4 | 5 | class Detection(object): 6 | """ 7 | This class represents a bounding box detection in a single image. 8 | 9 | Parameters 10 | ---------- 11 | tlwh : array_like 12 | Bounding box in format `(x, y, w, h)`. 13 | confidence : float 14 | Detector confidence score. 15 | feature : array_like 16 | A feature vector that describes the object contained in this image. 17 | 18 | Attributes 19 | ---------- 20 | tlwh : ndarray 21 | Bounding box in format `(top left x, top left y, width, height)`. 22 | confidence : ndarray 23 | Detector confidence score. 24 | feature : ndarray | NoneType 25 | A feature vector that describes the object contained in this image. 26 | 27 | """ 28 | 29 | def __init__(self, tlwh, rbox, confidence, feature): 30 | self.tlwh = np.asarray(tlwh, dtype=np.float32) 31 | self.rbox = np.asarray(rbox, dtype=np.float32) 32 | self.confidence = float(confidence) 33 | self.feature = np.asarray(feature.cpu(), dtype=np.float32) 34 | 35 | def to_tlbr(self): 36 | """Convert bounding box to format `(min x, min y, max x, max y)`, i.e., 37 | `(top left, bottom right)`. 38 | """ 39 | ret = self.tlwh.copy() 40 | ret[2:] += ret[:2] 41 | return ret 42 | 43 | def to_xyah(self): 44 | """Convert bounding box to format `(center x, center y, aspect ratio, 45 | height)`, where the aspect ratio is `width / height`. 46 | """ 47 | ret = self.tlwh.copy() 48 | ret[:2] += ret[2:] / 2 49 | ret[2] /= ret[3] 50 | return ret 51 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/sort/preprocessing.py: -------------------------------------------------------------------------------- 1 | # vim: expandtab:ts=4:sw=4 2 | import numpy as np 3 | import cv2 4 | 5 | 6 | def non_max_suppression(boxes, max_bbox_overlap, scores=None): 7 | """Suppress overlapping detections. 8 | 9 | Original code from [1]_ has been adapted to include confidence score. 10 | 11 | .. [1] http://www.pyimagesearch.com/2015/02/16/ 12 | faster-non-maximum-suppression-python/ 13 | 14 | Examples 15 | -------- 16 | 17 | >>> boxes = [d.roi for d in detections] 18 | >>> scores = [d.confidence for d in detections] 19 | >>> indices = non_max_suppression(boxes, max_bbox_overlap, scores) 20 | >>> detections = [detections[i] for i in indices] 21 | 22 | Parameters 23 | ---------- 24 | boxes : ndarray 25 | Array of ROIs (x, y, width, height). 26 | max_bbox_overlap : float 27 | ROIs that overlap more than this values are suppressed. 28 | scores : Optional[array_like] 29 | Detector confidence score. 30 | 31 | Returns 32 | ------- 33 | List[int] 34 | Returns indices of detections that have survived non-maxima suppression. 35 | 36 | """ 37 | if len(boxes) == 0: 38 | return [] 39 | 40 | boxes = boxes.astype(np.float) 41 | pick = [] 42 | 43 | x1 = boxes[:, 0] 44 | y1 = boxes[:, 1] 45 | x2 = boxes[:, 2] + boxes[:, 0] 46 | y2 = boxes[:, 3] + boxes[:, 1] 47 | 48 | area = (x2 - x1 + 1) * (y2 - y1 + 1) 49 | if scores is not None: 50 | idxs = np.argsort(scores) 51 | else: 52 | idxs = np.argsort(y2) 53 | 54 | while len(idxs) > 0: 55 | last = len(idxs) - 1 56 | i = idxs[last] 57 | pick.append(i) 58 | 59 | xx1 = np.maximum(x1[i], x1[idxs[:last]]) 60 | yy1 = np.maximum(y1[i], y1[idxs[:last]]) 61 | xx2 = np.minimum(x2[i], x2[idxs[:last]]) 62 | yy2 = np.minimum(y2[i], y2[idxs[:last]]) 63 | 64 | w = np.maximum(0, xx2 - xx1 + 1) 65 | h = np.maximum(0, yy2 - yy1 + 1) 66 | 67 | overlap = (w * h) / area[idxs[:last]] 68 | 69 | idxs = np.delete( 70 | idxs, np.concatenate( 71 | ([last], np.where(overlap > max_bbox_overlap)[0]))) 72 | 73 | return pick 74 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/YoloObbTrack/strong_sort/utils/__init__.py -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/asserts.py: -------------------------------------------------------------------------------- 1 | from os import environ 2 | 3 | 4 | def assert_in(file, files_to_check): 5 | if file not in files_to_check: 6 | raise AssertionError("{} does not exist in the list".format(str(file))) 7 | return True 8 | 9 | 10 | def assert_in_env(check_list: list): 11 | for item in check_list: 12 | assert_in(item, environ.keys()) 13 | return True 14 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/draw.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1) 5 | 6 | 7 | def compute_color_for_labels(label): 8 | """ 9 | Simple function that adds fixed color depending on the class 10 | """ 11 | color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette] 12 | return tuple(color) 13 | 14 | 15 | def draw_boxes(img, bbox, identities=None, offset=(0,0)): 16 | for i,box in enumerate(bbox): 17 | x1,y1,x2,y2 = [int(i) for i in box] 18 | x1 += offset[0] 19 | x2 += offset[0] 20 | y1 += offset[1] 21 | y2 += offset[1] 22 | # box text and bar 23 | id = int(identities[i]) if identities is not None else 0 24 | color = compute_color_for_labels(id) 25 | label = '{}{:d}'.format("", id) 26 | t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0] 27 | cv2.rectangle(img,(x1, y1),(x2,y2),color,3) 28 | cv2.rectangle(img,(x1, y1),(x1+t_size[0]+3,y1+t_size[1]+4), color,-1) 29 | cv2.putText(img,label,(x1,y1+t_size[1]+4), cv2.FONT_HERSHEY_PLAIN, 2, [255,255,255], 2) 30 | return img 31 | 32 | 33 | 34 | if __name__ == '__main__': 35 | for i in range(82): 36 | print(compute_color_for_labels(i)) 37 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/evaluation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import copy 4 | import motmetrics as mm 5 | mm.lap.default_solver = 'lap' 6 | from utils.io import read_results, unzip_objs 7 | 8 | 9 | class Evaluator(object): 10 | 11 | def __init__(self, data_root, seq_name, data_type): 12 | self.data_root = data_root 13 | self.seq_name = seq_name 14 | self.data_type = data_type 15 | 16 | self.load_annotations() 17 | self.reset_accumulator() 18 | 19 | def load_annotations(self): 20 | assert self.data_type == 'mot' 21 | 22 | gt_filename = os.path.join(self.data_root, self.seq_name, 'gt', 'gt.txt') 23 | self.gt_frame_dict = read_results(gt_filename, self.data_type, is_gt=True) 24 | self.gt_ignore_frame_dict = read_results(gt_filename, self.data_type, is_ignore=True) 25 | 26 | def reset_accumulator(self): 27 | self.acc = mm.MOTAccumulator(auto_id=True) 28 | 29 | def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False): 30 | # results 31 | trk_tlwhs = np.copy(trk_tlwhs) 32 | trk_ids = np.copy(trk_ids) 33 | 34 | # gts 35 | gt_objs = self.gt_frame_dict.get(frame_id, []) 36 | gt_tlwhs, gt_ids = unzip_objs(gt_objs)[:2] 37 | 38 | # ignore boxes 39 | ignore_objs = self.gt_ignore_frame_dict.get(frame_id, []) 40 | ignore_tlwhs = unzip_objs(ignore_objs)[0] 41 | 42 | 43 | # remove ignored results 44 | keep = np.ones(len(trk_tlwhs), dtype=bool) 45 | iou_distance = mm.distances.iou_matrix(ignore_tlwhs, trk_tlwhs, max_iou=0.5) 46 | if len(iou_distance) > 0: 47 | match_is, match_js = mm.lap.linear_sum_assignment(iou_distance) 48 | match_is, match_js = map(lambda a: np.asarray(a, dtype=int), [match_is, match_js]) 49 | match_ious = iou_distance[match_is, match_js] 50 | 51 | match_js = np.asarray(match_js, dtype=int) 52 | match_js = match_js[np.logical_not(np.isnan(match_ious))] 53 | keep[match_js] = False 54 | trk_tlwhs = trk_tlwhs[keep] 55 | trk_ids = trk_ids[keep] 56 | 57 | # get distance matrix 58 | iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5) 59 | 60 | # acc 61 | self.acc.update(gt_ids, trk_ids, iou_distance) 62 | 63 | if rtn_events and iou_distance.size > 0 and hasattr(self.acc, 'last_mot_events'): 64 | events = self.acc.last_mot_events # only supported by https://github.com/longcw/py-motmetrics 65 | else: 66 | events = None 67 | return events 68 | 69 | def eval_file(self, filename): 70 | self.reset_accumulator() 71 | 72 | result_frame_dict = read_results(filename, self.data_type, is_gt=False) 73 | frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys()))) 74 | for frame_id in frames: 75 | trk_objs = result_frame_dict.get(frame_id, []) 76 | trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2] 77 | self.eval_frame(frame_id, trk_tlwhs, trk_ids, rtn_events=False) 78 | 79 | return self.acc 80 | 81 | @staticmethod 82 | def get_summary(accs, names, metrics=('mota', 'num_switches', 'idp', 'idr', 'idf1', 'precision', 'recall')): 83 | names = copy.deepcopy(names) 84 | if metrics is None: 85 | metrics = mm.metrics.motchallenge_metrics 86 | metrics = copy.deepcopy(metrics) 87 | 88 | mh = mm.metrics.create() 89 | summary = mh.compute_many( 90 | accs, 91 | metrics=metrics, 92 | names=names, 93 | generate_overall=True 94 | ) 95 | 96 | return summary 97 | 98 | @staticmethod 99 | def save_summary(summary, filename): 100 | import pandas as pd 101 | writer = pd.ExcelWriter(filename) 102 | summary.to_excel(writer) 103 | writer.save() 104 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/log.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | 4 | def get_logger(name='root'): 5 | formatter = logging.Formatter( 6 | # fmt='%(asctime)s [%(levelname)s]: %(filename)s(%(funcName)s:%(lineno)s) >> %(message)s') 7 | fmt='%(asctime)s [%(levelname)s]: %(message)s', datefmt='%Y-%m-%d %H:%M:%S') 8 | 9 | handler = logging.StreamHandler() 10 | handler.setFormatter(formatter) 11 | 12 | logger = logging.getLogger(name) 13 | logger.setLevel(logging.INFO) 14 | logger.addHandler(handler) 15 | return logger 16 | 17 | 18 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/parser.py: -------------------------------------------------------------------------------- 1 | import os 2 | import yaml 3 | from easydict import EasyDict as edict 4 | 5 | 6 | class YamlParser(edict): 7 | """ 8 | This is yaml parser based on EasyDict. 9 | """ 10 | 11 | def __init__(self, cfg_dict=None, config_file=None): 12 | if cfg_dict is None: 13 | cfg_dict = {} 14 | 15 | if config_file is not None: 16 | assert(os.path.isfile(config_file)) 17 | with open(config_file, 'r') as fo: 18 | yaml_ = yaml.load(fo.read(), Loader=yaml.FullLoader) 19 | cfg_dict.update(yaml_) 20 | 21 | super(YamlParser, self).__init__(cfg_dict) 22 | 23 | def merge_from_file(self, config_file): 24 | with open(config_file, 'r') as fo: 25 | yaml_ = yaml.load(fo.read(), Loader=yaml.FullLoader) 26 | self.update(yaml_) 27 | 28 | def merge_from_dict(self, config_dict): 29 | self.update(config_dict) 30 | 31 | 32 | def get_config(config_file=None): 33 | return YamlParser(config_file=config_file) 34 | 35 | 36 | if __name__ == "__main__": 37 | cfg = YamlParser(config_file="../configs/yolov3.yaml") 38 | cfg.merge_from_file("../configs/strong_sort.yaml") 39 | 40 | import ipdb 41 | ipdb.set_trace() 42 | -------------------------------------------------------------------------------- /YoloObbTrack/strong_sort/utils/tools.py: -------------------------------------------------------------------------------- 1 | from functools import wraps 2 | from time import time 3 | 4 | 5 | def is_video(ext: str): 6 | """ 7 | Returns true if ext exists in 8 | allowed_exts for video files. 9 | 10 | Args: 11 | ext: 12 | 13 | Returns: 14 | 15 | """ 16 | 17 | allowed_exts = ('.mp4', '.webm', '.ogg', '.avi', '.wmv', '.mkv', '.3gp') 18 | return any((ext.endswith(x) for x in allowed_exts)) 19 | 20 | 21 | def tik_tok(func): 22 | """ 23 | keep track of time for each process. 24 | Args: 25 | func: 26 | 27 | Returns: 28 | 29 | """ 30 | @wraps(func) 31 | def _time_it(*args, **kwargs): 32 | start = time() 33 | try: 34 | return func(*args, **kwargs) 35 | finally: 36 | end_ = time() 37 | print("time: {:.03f}s, fps: {:.03f}".format(end_ - start, 1 / (end_ - start))) 38 | 39 | return _time_it 40 | -------------------------------------------------------------------------------- /__pycache__/val.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/__pycache__/val.cpython-37.pyc -------------------------------------------------------------------------------- /__pycache__/val.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/__pycache__/val.cpython-38.pyc -------------------------------------------------------------------------------- /data/dotav1_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 | 5 | # # path: /home/yuanzhengqian/yolov5_obb-master/dataset/jyz_outer_boundary # dataset root dir 6 | # train: /home/yuanzhengqian/yolov5_obb-master/dataset/jyz_outer_boundary/images #images # train images (relative to 'path') 7 | # val: /home/yuanzhengqian/yolov5_obb-master/dataset/jyz_outer_boundary/images #images # val images (relative to 'path') 8 | # # test: val/images #images # test images (optional) 9 | 10 | # # Classes 11 | # nc: 1 # number of classes 12 | # names: ['jyz_outer_boundary','jyz'] # class names 13 | 14 | 15 | 16 | 17 | # path: /home/yuanzhengqian/yolov5_obb-master/dataset/DOTA # dataset root dir 18 | # train: train/images # train images (relative to 'path') 19 | # val: val/images # val images (relative to 'path') 20 | # test: test/images # test images (optional) 21 | 22 | path: ./dataset/DOTA # dataset root dir 23 | train: train/images/train.txt #images # train images (relative to 'path') 24 | val: val/images/val.txt #images # val images (relative to 'path') 25 | test: val/images/val.txt #images # test images (optional) 26 | 27 | # Classes 28 | # nc: 16 # number of classes 29 | # names: ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 30 | # 'large-vehicle', 'ship', 'tennis-court', 'basketball-court', 'storage-tank', 31 | # 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter'] # class names 32 | nc: 16 # number of classes 33 | names: ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 34 | 'large-vehicle', 'ship', 'tennis-court', 'basketball-court', 'storage-tank', 35 | 'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter', 36 | 'container-crane'] # class names 37 | 38 | # Download script/URL (optional) 39 | # download: https://ultralytics.com/assets/coco128.zip 40 | -------------------------------------------------------------------------------- /data/hyps/obb/hyp.finetune_DroneVehicle.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.0032 4 | lrf: 0.12 5 | momentum: 0.843 6 | weight_decay: 0.00036 7 | warmup_epochs: 2.0 8 | warmup_momentum: 0.5 9 | warmup_bias_lr: 0.05 10 | box: 0.05 11 | cls: 0.5 12 | cls_pw: 1.0 13 | theta: 0.5 14 | theta_pw: 1.0 15 | obj: 1.0 16 | obj_pw: 1.0 17 | iou_t: 0.2 18 | anchor_t: 4.0 19 | fl_gamma: 0.0 20 | hsv_h: 0.015 21 | hsv_s: 0.7 22 | hsv_v: 0.4 23 | degrees: 0.0 24 | translate: 0.1 25 | scale: 0.25 26 | shear: 0.0 27 | perspective: 0.0 28 | flipud: 0.5 29 | fliplr: 0.5 30 | mosaic: 0.75 31 | mixup: 0.1 32 | copy_paste: 0.0 33 | cls_theta: 180 34 | csl_radius: 2.0 35 | 36 | -------------------------------------------------------------------------------- /data/hyps/obb/hyp.finetune_dota.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.001 4 | lrf: 0.2 5 | momentum: 0.937 6 | weight_decay: 0.0005 7 | warmup_epochs: 3.0 8 | warmup_momentum: 0.8 9 | warmup_bias_lr: 0.1 10 | box: 0.05 11 | cls: 0.5 12 | dfl: 0.05 13 | cls_pw: 1.0 14 | theta: 0.5 15 | theta_pw: 1.0 16 | obj: 1.0 17 | obj_pw: 1.0 18 | iou_t: 0.2 19 | anchor_t: 4.0 20 | fl_gamma: 1.5 21 | hsv_h: 0.015 22 | hsv_s: 0.7 23 | hsv_v: 0.4 24 | degrees: 0 25 | translate: 0.1 26 | scale: 0.25 27 | shear: 0.0 28 | perspective: 0.0 29 | flipud: 0.5 30 | fliplr: 0.5 31 | mosaic: 0.75 32 | # mosaic: 0 33 | mixup: 0.0 34 | copy_paste: 0.0 35 | cls_theta: 180 36 | csl_radius: 2.0 37 | no_rotato_ratio: 1.0 38 | use_kpt: False 39 | -------------------------------------------------------------------------------- /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: 0.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 -------------------------------------------------------------------------------- /data/hyps/obb/hyp.finetune_dota_kpt.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.001 4 | lrf: 0.2 5 | momentum: 0.937 6 | weight_decay: 0.0005 7 | warmup_epochs: 3.0 8 | warmup_momentum: 0.8 9 | warmup_bias_lr: 0.1 10 | box: 0.01 11 | dfl: 0.05 12 | cls: 0.5 13 | cls_pw: 1.0 14 | kpt: 18.0 15 | kobj: 1.0 16 | theta: 0.5 17 | theta_pw: 1.0 18 | obj: 1.0 19 | obj_pw: 1.0 20 | iou_t: 0.2 21 | anchor_t: 4.0 22 | fl_gamma: 1.5 23 | hsv_h: 0.015 24 | hsv_s: 0.7 25 | hsv_v: 0.4 26 | degrees: 0 27 | translate: 0.1 28 | scale: 0.25 29 | shear: 0.0 30 | perspective: 0.0 31 | flipud: 0.5 32 | fliplr: 0.5 33 | mosaic: 0.75 34 | # mosaic: 0 35 | mixup: 0.0 36 | copy_paste: 0.0 37 | cls_theta: 180 38 | csl_radius: 2.0 39 | no_rotato_ratio: 1.0 40 | use_kpt: True 41 | -------------------------------------------------------------------------------- /data/hyps/obb/hyp.paper.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.01 4 | lrf: 0.2 5 | momentum: 0.937 6 | weight_decay: 0.0005 7 | warmup_epochs: 3.0 8 | warmup_momentum: 0.8 9 | warmup_bias_lr: 0.1 10 | box: 0.05 11 | cls: 0.5 12 | cls_pw: 1.0 13 | theta: 0.5 14 | theta_pw: 1.0 15 | obj: 1.0 16 | obj_pw: 1.0 17 | iou_t: 0.2 18 | anchor_t: 4.0 19 | # anchors: 3.2 20 | fl_gamma: 0.0 21 | hsv_h: 0.0 22 | hsv_s: 0.0 23 | hsv_v: 0.0 24 | degrees: 0.0 25 | translate: 0.0 26 | scale: 0.5 27 | shear: 0.0 28 | perspective: 0.0 29 | flipud: 0.0 30 | fliplr: 0.0 31 | mosaic: 1.0 32 | mixup: 1.0 33 | copy_paste: 0.0 34 | cls_theta: 180 35 | csl_radius: 2.0 36 | -------------------------------------------------------------------------------- /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 - <>> import torch 30 | >>> torch.version.cuda 31 | >>> exit() 32 | ``` 33 | d. Clone the Yolov8_obb_Prune_Track repository. 34 | ``` 35 | git clone https://github.com/yzqxy/Yolov8_obb_Prune_Track.git 36 | cd Yolov8_obb_Prune_Track 37 | ``` 38 | e. Install Yolov8_obb_Prune_Track. 39 | ```python 40 | pip install -r requirements.txt 41 | ``` 42 | 43 | ## Install DOTA_devkit. 44 | **(Custom Install, it's just a tool to split the high resolution image and evaluation the obb)** 45 | ``` 46 | cd Yolov8_obb_Prune_Track/DOTA_devkit 47 | sudo apt-get install swig 48 | swig -c++ -python polyiou.i 49 | python setup.py build_ext --inplace 50 | ``` 51 | 52 | -------------------------------------------------------------------------------- /docs/obb_kpt_result.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/docs/obb_kpt_result.jpg -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__init__.py -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/common.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/common.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/common.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/common.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/common_prune.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/common_prune.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/common_prune.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/common_prune.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/experimental.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/experimental.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/experimental.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/experimental.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/yolo.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/yolo.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/yolo.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/yolo.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/yolo_prune.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/yolo_prune.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/yolo_prune.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/models/__pycache__/yolo_prune.cpython-38.pyc -------------------------------------------------------------------------------- /models/hub/yolov8n_C2FTS.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2fTR, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C2fTR, [1024]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | 44 | # head: 45 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 46 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 47 | # - [-1, 3, C2f, [512]] # 13 48 | 49 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 50 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 51 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 52 | 53 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 54 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 55 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 56 | 57 | # - [-1, 1, Conv, [256, 3, 2]] 58 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 59 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 60 | 61 | # - [-1, 1, Conv, [256, 3, 2]] 62 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 63 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 64 | 65 | # - [-1, 1, Conv, [512, 3, 2]] 66 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 67 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 68 | 69 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/hub/yolov8n_C3TS.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C3TR, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C3TR, [1024]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | 44 | # head: 45 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 46 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 47 | # - [-1, 3, C2f, [512]] # 13 48 | 49 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 50 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 51 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 52 | 53 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 54 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 55 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 56 | 57 | # - [-1, 1, Conv, [256, 3, 2]] 58 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 59 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 60 | 61 | # - [-1, 1, Conv, [256, 3, 2]] 62 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 63 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 64 | 65 | # - [-1, 1, Conv, [512, 3, 2]] 66 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 67 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 68 | 69 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/hub/yolov8n_RFCAConv.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | [-1, 1, RFCAConv2, [512, 1, 1]], 37 | 38 | [-1, 1, Conv, [512, 3, 2]], 39 | [[-1, 9], 1, Concat, [1]], # cat head P5 40 | [-1, 3, C2f, [1024]], # 23 (P5/32-large) 41 | 42 | [[15, 18, 22], 1, Detect_v8, [nc]], 43 | ] # Detect(P3, P4, P5) 44 | -------------------------------------------------------------------------------- /models/hub/yolov8n_c2fTS_cbam.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2fTR, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | [-1, 1, CBAM, [512]], 37 | 38 | [-1, 1, Conv, [512, 3, 2]], 39 | [[-1, 9], 1, Concat, [1]], # cat head P5 40 | [-1, 3, C2fTR, [1024]], # 23 (P5/32-large) 41 | [-1, 1, CBAM, [1024]], 42 | 43 | [[15, 18, 22], 1, Detect_v8, [nc]], 44 | ] # Detect(P3, P4, P5) 45 | 46 | # head: 47 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 48 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 49 | # - [-1, 3, C2f, [512]] # 13 50 | 51 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 52 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 53 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 54 | 55 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 56 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 57 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 58 | 59 | # - [-1, 1, Conv, [256, 3, 2]] 60 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 61 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 62 | 63 | # - [-1, 1, Conv, [256, 3, 2]] 64 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 65 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 66 | 67 | # - [-1, 1, Conv, [512, 3, 2]] 68 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 69 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 70 | 71 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/hub/yolov8n_ca.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 1 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | [-1, 1, CAConv, [1024, 3,2]], 22 | ] # 9 23 | 24 | # YOLOv8.0n head 25 | head: 26 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 27 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 28 | [-1, 3, C2f, [512]], # 13 29 | 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 32 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 33 | 34 | [-1, 1, Conv, [256, 3, 2]], 35 | [[-1, 13,6], 1, Concat, [1]], # cat head P4 36 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 37 | 38 | [-1, 1, Conv, [512, 3, 2]], 39 | [[-1, 9], 1, Concat, [1]], # cat head P5 40 | [-1, 3, C2f, [1024]], # 23 (P5/32-large) 41 | 42 | [[16, 19, 22], 1, Detect_v8, [nc]], 43 | ] # Detect(P3, P4, P5) 44 | 45 | # head: 46 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 47 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 48 | # - [-1, 3, C2f, [512]] # 13 49 | 50 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 51 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 52 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 53 | 54 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 55 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 56 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 57 | 58 | # - [-1, 1, Conv, [256, 3, 2]] 59 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 60 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 61 | 62 | # - [-1, 1, Conv, [256, 3, 2]] 63 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 64 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 65 | 66 | # - [-1, 1, Conv, [512, 3, 2]] 67 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 68 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 69 | 70 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/hub/yolov8n_cbam.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | [-1, 1, CBAM, [512]], 37 | 38 | [-1, 1, Conv, [512, 3, 2]], 39 | [[-1, 9], 1, Concat, [1]], # cat head P5 40 | [-1, 3, C2f, [1024]], # 23 (P5/32-large) 41 | [-1, 1, CBAM, [1024]], 42 | 43 | [[15, 18, 23], 1, Detect_v8, [nc]], 44 | ] # Detect(P3, P4, P5) 45 | 46 | # head: 47 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 48 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 49 | # - [-1, 3, C2f, [512]] # 13 50 | 51 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 52 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 53 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 54 | 55 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 56 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 57 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 58 | 59 | # - [-1, 1, Conv, [256, 3, 2]] 60 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 61 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 62 | 63 | # - [-1, 1, Conv, [256, 3, 2]] 64 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 65 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 66 | 67 | # - [-1, 1, Conv, [512, 3, 2]] 68 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 69 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 70 | 71 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/yaml/yolov8l.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 1.00 # scales module repeats 6 | width_multiple: 1.00 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [512, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [512, True]], 20 | [-1, 1, SPPF, [512, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C2f, [512]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | 44 | -------------------------------------------------------------------------------- /models/yaml/yolov8n.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C2f, [1024]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | 44 | # head: 45 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 46 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 47 | # - [-1, 3, C2f, [512]] # 13 48 | 49 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 50 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 51 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 52 | 53 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 54 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 55 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 56 | 57 | # - [-1, 1, Conv, [256, 3, 2]] 58 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 59 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 60 | 61 | # - [-1, 1, Conv, [256, 3, 2]] 62 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 63 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 64 | 65 | # - [-1, 1, Conv, [512, 3, 2]] 66 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 67 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 68 | 69 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/yaml/yolov8n_kpt.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C2f, [1024]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8_kpt, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | -------------------------------------------------------------------------------- /models/yaml/yolov8nx6.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 26 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 27 | - [-1, 3, C2f, [512]] # 13 28 | 29 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 30 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 31 | - [-1, 3, C2f, [512]] # 17 (P3/8-small) 32 | 33 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 34 | - [[-1, 2], 1, Concat, [1]] # cat backbone P3 35 | - [-1, 3, C2f, [256]] # 20 (P3/8-small) 36 | 37 | - [-1, 1, Conv, [256, 3, 2]] 38 | - [[-1, 15], 1, Concat, [1]] # cat head P4 39 | - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 40 | 41 | - [-1, 1, Conv, [256, 3, 2]] 42 | - [[-1, 12], 1, Concat, [1]] # cat head P4 43 | - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 44 | 45 | - [-1, 1, Conv, [512, 3, 2]] 46 | - [[-1, 9], 1, Concat, [1]] # cat head P5 47 | - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 48 | 49 | - [[18, 21, 24, 27], 1, Detect_v8, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/yaml/yolov8s.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 1 # number of classes 5 | depth_multiple: 0.33 # scales module repeats 6 | width_multiple: 0.50 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [1024, True]], 20 | [-1, 1, SPPF, [1024, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C2f, [1024]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | 44 | # head: 45 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 46 | # - [[-1, 6], 1, Concat, [1]] # cat backbone P4 47 | # - [-1, 3, C2f, [512]] # 13 48 | 49 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 50 | # - [[-1, 4], 1, Concat, [1]] # cat backbone P3 51 | # - [-1, 3, C2f, [512]] # 17 (P3/8-small) 52 | 53 | # - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 54 | # - [[-1, 2], 1, Concat, [1]] # cat backbone P3 55 | # - [-1, 3, C2f, [256]] # 20 (P3/8-small) 56 | 57 | # - [-1, 1, Conv, [256, 3, 2]] 58 | # - [[-1, 15], 1, Concat, [1]] # cat head P4 59 | # - [-1, 3, C2f, [256]] # 23 (P4/16-medium) 60 | 61 | # - [-1, 1, Conv, [256, 3, 2]] 62 | # - [[-1, 12], 1, Concat, [1]] # cat head P4 63 | # - [-1, 3, C2f, [512]] # 26 (P4/16-medium) 64 | 65 | # - [-1, 1, Conv, [512, 3, 2]] 66 | # - [[-1, 9], 1, Concat, [1]] # cat head P5 67 | # - [-1, 3, C2f, [1024]] # 29 (P5/32-large) 68 | 69 | # - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /models/yaml/yolov8x.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 2 # number of classes 5 | depth_multiple: 1.00 # scales module repeats 6 | width_multiple: 1.25 # scales convolution channels 7 | 8 | # YOLOv8.0n backbone 9 | backbone: 10 | # [from, repeats, module, args] 11 | [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 12 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 13 | [-1, 3, C2f, [128, True]], 14 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 15 | [-1, 6, C2f, [256, True]], 16 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 17 | [-1, 6, C2f, [512, True]], 18 | [-1, 1, Conv, [512, 3, 2]], # 7-P5/32 19 | [-1, 3, C2f, [512, True]], 20 | [-1, 1, SPPF, [512, 5]], 21 | ] # 9 22 | 23 | # YOLOv8.0n head 24 | head: 25 | [[-1, 1, nn.Upsample, [None, 2, 'nearest']], 26 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 27 | [-1, 3, C2f, [512]], # 13 28 | 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 31 | [-1, 3, C2f, [256]], # 17 (P3/8-small) 32 | 33 | [-1, 1, Conv, [256, 3, 2]], 34 | [[-1, 12], 1, Concat, [1]], # cat head P4 35 | [-1, 3, C2f, [512]], # 20 (P4/16-medium) 36 | 37 | [-1, 1, Conv, [512, 3, 2]], 38 | [[-1, 9], 1, Concat, [1]], # cat head P5 39 | [-1, 3, C2f, [512]], # 23 (P5/32-large) 40 | 41 | [[15, 18, 21], 1, Detect_v8, [nc]], 42 | ] # Detect(P3, P4, P5) 43 | 44 | -------------------------------------------------------------------------------- /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 | opencv-python==4.7.0 8 | Pillow>=7.1.2 9 | PyYAML>=5.3.1 10 | requests>=2.23.0 11 | scipy>=1.4.1 12 | # torch>=1.7.0# 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 | onnx==1.13.0 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 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | # Project-wide configuration file, can be used for package metadata and other toll configurations 2 | # Example usage: global configuration for PEP8 (via flake8) setting or default pytest arguments 3 | 4 | [metadata] 5 | license_file = LICENSE 6 | description-file = README.md 7 | 8 | 9 | [tool:pytest] 10 | norecursedirs = 11 | .git 12 | dist 13 | build 14 | addopts = 15 | --doctest-modules 16 | --durations=25 17 | --color=yes 18 | 19 | 20 | [flake8] 21 | max-line-length = 120 22 | exclude = .tox,*.egg,build,temp 23 | select = E,W,F 24 | doctests = True 25 | verbose = 2 26 | # https://pep8.readthedocs.io/en/latest/intro.html#error-codes 27 | format = pylint 28 | # see: https://www.flake8rules.com/ 29 | ignore = 30 | E731 # Do not assign a lambda expression, use a def 31 | F405 32 | E402 33 | F841 34 | E741 35 | F821 36 | E722 37 | F401 38 | W504 39 | E127 40 | W504 41 | E231 42 | E501 43 | F403 44 | E302 45 | F541 46 | 47 | 48 | [isort] 49 | # https://pycqa.github.io/isort/docs/configuration/options.html 50 | line_length = 120 51 | multi_line_output = 0 52 | -------------------------------------------------------------------------------- /sh/ddp_train.sh: -------------------------------------------------------------------------------- 1 | python -m torch.distributed.launch --nproc_per_node 3 train.py \ 2 | --batch 32 \ 3 | --data '/home/test/Persons/hukaixuan/yolov5/datasets/package/data.yaml' \ 4 | --weights 'runs/train/exp_multi_aug_yolov5n/weights/last.pt' \ 5 | --hyp 'data/hyps/hyp.finetune_package.yaml' \ 6 | --device 2 \ 7 | --epochs 50 \ 8 | --img 512 \ 9 | --cache \ 10 | --sync-bn 11 | 12 | 13 | tensorboard --logdir runs/train/exp 14 | 15 | # 测试 train/val/test 数据集的情况 16 | python val.py \ 17 | --weights runs/train/exp_close_mosaic/weights/last.pt \ 18 | --img 480 \ 19 | --half \ 20 | --batch 1 \ 21 | --conf 0.001 --iou 0.65 \ # mAP 测精度的标准 22 | --conf 0.25 --iou 0.45 \ # speed 测速标准 23 | --task 'val' or 'train' or 'test' 24 | 25 | # speed模式 26 | python val.py \ 27 | --weights runs/train/exp_close_mosaic/weights/best.pt runs/train/exp_close_mosaic/weights/last.pt \ 28 | --img 480 \ 29 | --half \ 30 | --task 'speed' \ 31 | --device 'cpu' \ 32 | --batch 1 33 | 34 | #mAP 35 | python val.py --task 'test' --batch 16 --save-json --name 'yolov5t_dotav1_test_split' 36 | python tools/TestJson2VocClassTxt.py \ 37 | --json_path 'runs/val/yolov5t_DroneVehicle_val/best_obb_predictions.json' \ 38 | --save_path 'runs/val/yolov5t_DroneVehicle_val/splited_obb_prediction_Txt' 39 | python DOTA_devkit/ResultMerge_multi_process.py \ 40 | --scrpath 'runs/val/yolov5t_dotav1_test_split/splited_obb_prediction_Txt' \ 41 | --dstpath 'runs/val/yolov5t_dotav1_test_split/Merged_obb_prediction_Txt' 42 | python DOTA_devkit/results_obb2hbb.py \ 43 | --srcpath 'runs/val/yolov5t_dotav1_test_split/Merged_obb_prediction_Txt' \ 44 | --dstpath 'runs/val/yolov5t_dotav1_test_split/Merged_hbb_prediction_Txt' 45 | -------------------------------------------------------------------------------- /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!") -------------------------------------------------------------------------------- /tools/Xml2Txt.py: -------------------------------------------------------------------------------- 1 | import xml.etree.ElementTree as ET 2 | from DOTA_devkit import dota_utils as util 3 | import os 4 | import shutil 5 | 6 | def DroneVehicle2DOTA(xmlpath, txtpath, extractclassname, specified_class): 7 | """ 8 | Trans DroneVehicle xml farmat to DOTA labels format 9 | Args: 10 | xmlpath: the path of xml with DroneVehicle2DOTA format 11 | txtpath: the path of txt with DOTA format 12 | extractclassname: the category 13 | specified_class: only output single category 14 | """ 15 | if os.path.exists(txtpath): 16 | shutil.rmtree(txtpath) # delete output folder 17 | os.makedirs(txtpath) # make new output folder 18 | filelist = util.GetFileFromThisRootDir(xmlpath) # fileist=['/.../0001.xml', ...] 19 | for xmlfile in filelist: # fullname='/.../0001.xml' 20 | name = os.path.splitext(os.path.basename(xmlfile))[0] # name='0001' 21 | out_file = open(os.path.join(txtpath, name + '.txt'), 'w') 22 | tree = ET.parse(xmlfile) 23 | root = tree.getroot() 24 | for obj in root.iter('object'): 25 | cls = obj.find('name') 26 | if cls == None: 27 | continue 28 | cls = cls.text 29 | diffcult = obj.find('difficult') 30 | diffcult = int(diffcult.text) if diffcult != None else 0 31 | 32 | if diffcult < 2: 33 | # cls = cls.replace(' ', '_') 34 | cls = specified_class 35 | polygon = obj.find('polygon') 36 | if polygon == None: 37 | continue 38 | polygon = [int(polygon.find(x).text) for x in ('x1', 'y1', 'x2', 'y2', 'x3', 'y3', 'x4', 'y4')] 39 | out_file.write(" ".join([str(a) for a in (*polygon, cls, diffcult)]) + '\n') 40 | else: 41 | print(f'{cls} is not in the extractclassname or diffcult is {diffcult}') 42 | 43 | if __name__ == "__main__": 44 | xmlpath = ['/media/test/DroneVehicle/val/raw/vallabel', 45 | '/media/test/DroneVehicle/val/rays/vallabelr', 46 | '/media/test/DroneVehicle/train/raw/trainlabel', 47 | '/media/test/DroneVehicle/train/rays/trainlabelr'] 48 | 49 | txtpath = ['/media/test/DroneVehicle/val/raw/vallabel_txt', 50 | '/media/test/DroneVehicle/val/rays/vallabelr_txt', 51 | '/media/test/DroneVehicle/train/raw/trainlabel_txt', 52 | '/media/test/DroneVehicle/train/rays/trainlabelr_txt'] 53 | extractclassname = {'car', 'truck', 'feright_car', 'feright car', 'bus', 'van'} 54 | specified_class = 'vehicle' 55 | for (xml, txt) in zip(xmlpath, txtpath): 56 | print(f"{xml} is converting to {txt}") 57 | DroneVehicle2DOTA(xml, txt, extractclassname, specified_class) -------------------------------------------------------------------------------- /tools/mk_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | def make_dataset(opt): 4 | data_path=opt.data_path 5 | with open(data_path+'/train.txt','w') as f: 6 | for filename1 in os.listdir(data_path): 7 | if filename1=='images': 8 | path2=data_path+'/'+filename1 9 | for filename2 in os.listdir(path2): 10 | path3=path2+'/'+filename2 11 | if filename2=='train': 12 | for filename3 in os.listdir(path3): 13 | f.write(path3+'/'+filename3+'\n') 14 | with open(data_path+'/val.txt','w') as f: 15 | for filename1 in os.listdir(data_path): 16 | if filename1=='images': 17 | path2=data_path+'/'+filename1 18 | for filename2 in os.listdir(path2): 19 | path3=path2+'/'+filename2 20 | if filename2=='val': 21 | for filename3 in os.listdir(path3): 22 | f.write(path3+'/'+filename3+'\n') 23 | 24 | def parse_opt(known=False): 25 | parser = argparse.ArgumentParser() 26 | parser.add_argument('--data_path', type=str, default='datasets/UAV') 27 | 28 | opt = parser.parse_known_args()[0] if known else parser.parse_args() 29 | return opt 30 | 31 | if __name__ == "__main__": 32 | opt = parse_opt() 33 | make_dataset(opt) -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/activations.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/activations.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/augmentations.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/augmentations.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/augmentations.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/augmentations.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/autoanchor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/autoanchor.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/autoanchor.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/autoanchor.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/autobatch.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/autobatch.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/autobatch.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/autobatch.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/callbacks.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/callbacks.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/callbacks.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/callbacks.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/datasets.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/datasets.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/datasets.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/datasets.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/downloads.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/downloads.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/downloads.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/downloads.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/general.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/general.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/general.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/general.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/lion_pytorch.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/lion_pytorch.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/lion_pytorch.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/lion_pytorch.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/loss.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/loss.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/metrics.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/metrics.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plots.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/plots.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plots.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/plots.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/prune_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/prune_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/prune_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/prune_utils.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/rboxs_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/rboxs_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/rboxs_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/rboxs_utils.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/tal.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/tal.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/tal.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/tal.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/torch_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/torch_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/torch_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/__pycache__/torch_utils.cpython-38.pyc -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/aws/__init__.py -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/lion_pytorch.py: -------------------------------------------------------------------------------- 1 | # You may obtain a copy of the License at 2 | # 3 | # http://www.apache.org/licenses/LICENSE-2.0 4 | # 5 | # Unless required by applicable law or agreed to in writing, software 6 | # distributed under the License is distributed on an "AS IS" BASIS, 7 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 8 | # See the License for the specific language governing permissions and 9 | # limitations under the License. 10 | # ============================================================================== 11 | """PyTorch implementation of the Lion optimizer.""" 12 | import torch 13 | from torch.optim.optimizer import Optimizer 14 | 15 | 16 | class Lion(Optimizer): 17 | r"""Implements Lion algorithm.""" 18 | 19 | def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0): 20 | """Initialize the hyperparameters. 21 | 22 | Args: 23 | params (iterable): iterable of parameters to optimize or dicts defining 24 | parameter groups 25 | lr (float, optional): learning rate (default: 1e-4) 26 | betas (Tuple[float, float], optional): coefficients used for computing 27 | running averages of gradient and its square (default: (0.9, 0.99)) 28 | weight_decay (float, optional): weight decay coefficient (default: 0) 29 | """ 30 | 31 | if not 0.0 <= lr: 32 | raise ValueError('Invalid learning rate: {}'.format(lr)) 33 | if not 0.0 <= betas[0] < 1.0: 34 | raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0])) 35 | if not 0.0 <= betas[1] < 1.0: 36 | raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1])) 37 | defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) 38 | super().__init__(params, defaults) 39 | 40 | @torch.no_grad() 41 | def step(self, closure=None): 42 | """Performs a single optimization step. 43 | 44 | Args: 45 | closure (callable, optional): A closure that reevaluates the model 46 | and returns the loss. 47 | 48 | Returns: 49 | the loss. 50 | """ 51 | loss = None 52 | if closure is not None: 53 | with torch.enable_grad(): 54 | loss = closure() 55 | 56 | for group in self.param_groups: 57 | for p in group['params']: 58 | if p.grad is None: 59 | continue 60 | 61 | # Perform stepweight decay 62 | p.data.mul_(1 - group['lr'] * group['weight_decay']) 63 | 64 | grad = p.grad 65 | state = self.state[p] 66 | # State initialization 67 | if len(state) == 0: 68 | # Exponential moving average of gradient values 69 | state['exp_avg'] = torch.zeros_like(p) 70 | 71 | exp_avg = state['exp_avg'] 72 | beta1, beta2 = group['betas'] 73 | 74 | # Weight update 75 | update = exp_avg * beta1 + grad * (1 - beta1) 76 | p.add_(torch.sign(update), alpha=-group['lr']) 77 | # Decay the momentum running average coefficient 78 | exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) 79 | 80 | return loss -------------------------------------------------------------------------------- /utils/loggers/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /utils/loggers/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /utils/loggers/wandb/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__init__.py -------------------------------------------------------------------------------- /utils/loggers/wandb/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /utils/loggers/wandb/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /utils/loggers/wandb/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /utils/loggers/wandb/__pycache__/wandb_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__pycache__/wandb_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/loggers/wandb/__pycache__/wandb_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__pycache__/wandb_utils.cpython-38.pyc -------------------------------------------------------------------------------- /utils/loggers/wandb/__pycache__/wandb_utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/loggers/wandb/__pycache__/wandb_utils.cpython-39.pyc -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/nms_rotated/__init__.py: -------------------------------------------------------------------------------- 1 | from .nms_rotated_wrapper import obb_nms, poly_nms 2 | 3 | __all__ = ['obb_nms', 'poly_nms'] 4 | -------------------------------------------------------------------------------- /utils/nms_rotated/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /utils/nms_rotated/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /utils/nms_rotated/__pycache__/nms_rotated_wrapper.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/__pycache__/nms_rotated_wrapper.cpython-37.pyc -------------------------------------------------------------------------------- /utils/nms_rotated/__pycache__/nms_rotated_wrapper.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/__pycache__/nms_rotated_wrapper.cpython-38.pyc -------------------------------------------------------------------------------- /utils/nms_rotated/build/lib.linux-x86_64-cpython-37/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/lib.linux-x86_64-cpython-37/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/nms_rotated_cpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/nms_rotated_cpu.o -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/nms_rotated_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/nms_rotated_cuda.o -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/nms_rotated_ext.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/nms_rotated_ext.o -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/poly_nms_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-37/src/poly_nms_cuda.o -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-38/src/nms_rotated_cpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-38/src/nms_rotated_cpu.o -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-38/src/nms_rotated_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-38/src/nms_rotated_cuda.o -------------------------------------------------------------------------------- /utils/nms_rotated/build/temp.linux-x86_64-cpython-38/src/nms_rotated_ext.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/build/temp.linux-x86_64-cpython-38/src/nms_rotated_ext.o -------------------------------------------------------------------------------- /utils/nms_rotated/nms_rotated.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: nms-rotated 3 | Version: 0.0.0 4 | -------------------------------------------------------------------------------- /utils/nms_rotated/nms_rotated.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | setup.py 2 | nms_rotated.egg-info/PKG-INFO 3 | nms_rotated.egg-info/SOURCES.txt 4 | nms_rotated.egg-info/dependency_links.txt 5 | nms_rotated.egg-info/not-zip-safe 6 | nms_rotated.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 -------------------------------------------------------------------------------- /utils/nms_rotated/nms_rotated.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /utils/nms_rotated/nms_rotated.egg-info/not-zip-safe: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /utils/nms_rotated/nms_rotated.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /utils/nms_rotated/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/utils/nms_rotated/nms_rotated_ext.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/nms_rotated/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import os 3 | import subprocess 4 | import time 5 | from setuptools import find_packages, setup 6 | 7 | import torch 8 | from torch.utils.cpp_extension import (BuildExtension, CppExtension, 9 | CUDAExtension) 10 | def make_cuda_ext(name, module, sources, sources_cuda=[]): 11 | 12 | define_macros = [] 13 | extra_compile_args = {'cxx': []} 14 | 15 | if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1': 16 | define_macros += [('WITH_CUDA', None)] 17 | extension = CUDAExtension 18 | extra_compile_args['nvcc'] = [ 19 | '-D__CUDA_NO_HALF_OPERATORS__', 20 | '-D__CUDA_NO_HALF_CONVERSIONS__', 21 | '-D__CUDA_NO_HALF2_OPERATORS__', 22 | ] 23 | sources += sources_cuda 24 | else: 25 | print(f'Compiling {name} without CUDA') 26 | extension = CppExtension 27 | # raise EnvironmentError('CUDA is required to compile MMDetection!') 28 | 29 | return extension( 30 | name=f'{module}.{name}', 31 | sources=[os.path.join(*module.split('.'), p) for p in sources], 32 | define_macros=define_macros, 33 | extra_compile_args=extra_compile_args) 34 | 35 | # python setup.py develop 36 | if __name__ == '__main__': 37 | #write_version_py() 38 | setup( 39 | name='nms_rotated', 40 | ext_modules=[ 41 | make_cuda_ext( 42 | name='nms_rotated_ext', 43 | module='', 44 | sources=[ 45 | 'src/nms_rotated_cpu.cpp', 46 | 'src/nms_rotated_ext.cpp' 47 | ], 48 | sources_cuda=[ 49 | 'src/nms_rotated_cuda.cu', 50 | 'src/poly_nms_cuda.cu', 51 | ]), 52 | ], 53 | cmdclass={'build_ext': BuildExtension}, 54 | zip_safe=False) -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /utils/prune_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Time : 2021/5/24 下午4:36 3 | # @Author : midaskong 4 | # @File : prune_utils.py 5 | # @Description: 6 | 7 | import torch 8 | from copy import deepcopy 9 | import numpy as np 10 | import torch.nn.functional as F 11 | 12 | 13 | def gather_bn_weights(module_list): 14 | prune_idx = list(range(len(module_list))) 15 | size_list = [idx.weight.data.shape[0] for idx in module_list.values()] 16 | bn_weights = torch.zeros(sum(size_list)) 17 | index = 0 18 | for i, idx in enumerate(module_list.values()): 19 | size = size_list[i] 20 | bn_weights[index:(index + size)] = idx.weight.data.abs().clone() 21 | index += size 22 | return bn_weights 23 | 24 | def gather_conv_weights(module_list): 25 | prune_idx = list(range(len(module_list))) 26 | size_list = [idx.weight.data.shape[0] for idx in module_list.values()] 27 | 28 | conv_weights = torch.zeros(sum(size_list)) 29 | index = 0 30 | for i, idx in enumerate(module_list.values()): 31 | size = size_list[i] 32 | conv_weights[index:(index + size)] = idx.weight.data.abs().sum(dim=1).sum(dim=1).sum(dim=1).clone() 33 | index += size 34 | return conv_weights 35 | 36 | 37 | def obtain_bn_mask(bn_module, thre): 38 | 39 | thre = thre.cuda() 40 | mask = bn_module.weight.data.abs().ge(thre).float() 41 | if int(mask.sum())%8==0: 42 | return mask 43 | else: 44 | x=0 45 | num=8-(int(mask.sum())%8) 46 | for i in range(len(mask)): 47 | if mask[i]==0: 48 | mask[i]=1 49 | x=x+1 50 | if x==num: 51 | break 52 | return mask 53 | 54 | 55 | def obtain_conv_mask(conv_module, thre): 56 | thre = thre.cuda() 57 | mask = conv_module.weight.data.abs().sum(dim=1).sum(dim=1).sum(dim=1).ge(thre).float() 58 | return mask 59 | 60 | def uodate_pruned_yolov5_cfg(model, maskbndict): 61 | # save pruned yolov5 model in yaml format: 62 | # model: 63 | # model to be pruned 64 | # maskbndict: 65 | # key : module name 66 | # value : bn layer mask index 67 | return -------------------------------------------------------------------------------- /weights/yolov8n.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yzqxy/Yolov8_obb_Prune_Track/fb8dae0315540527f04615b516329663ab511429/weights/yolov8n.pt --------------------------------------------------------------------------------