├── LICENSE.txt ├── README.md ├── Utils ├── .DS_Store ├── Initailize_model.py └── __pycache__ │ ├── Compute_FDR.cpython-39.pyc │ ├── Confusion_mats.cpython-39.pyc │ ├── Generate_Learning_Curves.cpython-39.pyc │ ├── Generate_TSNE_visual.cpython-39.pyc │ ├── Network_functions.cpython-311.pyc │ ├── Network_functions.cpython-39.pyc │ ├── Save_Results.cpython-39.pyc │ └── pytorchtools.cpython-39.pyc ├── demo.py ├── ultralytics ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-311.pyc │ └── __init__.cpython-39.pyc ├── architecture.png ├── assets │ ├── bus.jpg │ └── zidane.jpg ├── cfg │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ └── __init__.cpython-39.pyc │ ├── datasets │ │ ├── Argoverse.yaml │ │ ├── DOTAv2.yaml │ │ ├── GlobalWheat2020.yaml │ │ ├── ImageNet.yaml │ │ ├── Objects365.yaml │ │ ├── SKU-110K.yaml │ │ ├── VOC.yaml │ │ ├── VisDrone.yaml │ │ ├── coco-pose.yaml │ │ ├── coco.yaml │ │ ├── coco128-seg.yaml │ │ ├── coco128.yaml │ │ ├── coco8-pose.yaml │ │ ├── coco8-seg.yaml │ │ ├── coco8.yaml │ │ ├── open-images-v7.yaml │ │ ├── tiger-pose.yaml │ │ └── xView.yaml │ ├── default.yaml │ ├── models │ │ ├── README.md │ │ ├── rt-detr │ │ │ ├── rtdetr-l.yaml │ │ │ └── rtdetr-x.yaml │ │ ├── v3 │ │ │ ├── yolov3-spp.yaml │ │ │ ├── yolov3-tiny.yaml │ │ │ └── yolov3.yaml │ │ ├── v5 │ │ │ ├── yolov5-p6.yaml │ │ │ └── yolov5.yaml │ │ ├── v6 │ │ │ └── yolov6.yaml │ │ └── v8 │ │ │ ├── stnyolov8.yaml │ │ │ ├── test.yaml │ │ │ ├── yolov8-cls.yaml │ │ │ ├── yolov8-p2.yaml │ │ │ ├── yolov8-p6.yaml │ │ │ ├── yolov8-pose-p6.yaml │ │ │ ├── yolov8-pose.yaml │ │ │ ├── yolov8-rtdetr.yaml │ │ │ ├── yolov8-seg-p6.yaml │ │ │ ├── yolov8-seg.yaml │ │ │ └── yolov8.yaml │ └── trackers │ │ ├── botsort.yaml │ │ └── bytetrack.yaml ├── data │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── augment.cpython-311.pyc │ │ ├── augment.cpython-39.pyc │ │ ├── base.cpython-311.pyc │ │ ├── base.cpython-39.pyc │ │ ├── build.cpython-311.pyc │ │ ├── build.cpython-39.pyc │ │ ├── converter.cpython-311.pyc │ │ ├── converter.cpython-39.pyc │ │ ├── dataset.cpython-311.pyc │ │ ├── dataset.cpython-39.pyc │ │ ├── loaders.cpython-311.pyc │ │ ├── loaders.cpython-39.pyc │ │ ├── utils.cpython-311.pyc │ │ └── utils.cpython-39.pyc │ ├── annotator.py │ ├── augment.py │ ├── base.py │ ├── build.py │ ├── converter.py │ ├── dataset.py │ ├── loaders.py │ ├── scripts │ │ ├── download_weights.sh │ │ ├── get_coco.sh │ │ ├── get_coco128.sh │ │ └── get_imagenet.sh │ └── utils.py ├── engine │ ├── __pycache__ │ │ ├── __init__.cpython-39.pyc │ │ ├── exporter.cpython-311.pyc │ │ ├── exporter.cpython-39.pyc │ │ ├── model.cpython-311.pyc │ │ ├── model.cpython-39.pyc │ │ ├── predictor.cpython-311.pyc │ │ ├── predictor.cpython-39.pyc │ │ ├── results.cpython-311.pyc │ │ ├── results.cpython-39.pyc │ │ ├── trainer.cpython-311.pyc │ │ ├── trainer.cpython-39.pyc │ │ ├── validator.cpython-311.pyc │ │ └── validator.cpython-39.pyc │ ├── exporter.py │ ├── model.py │ ├── predictor.py │ ├── results.py │ ├── trainer.py │ ├── tuner.py │ └── validator.py ├── hub │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── auth.cpython-311.pyc │ │ ├── auth.cpython-39.pyc │ │ ├── utils.cpython-311.pyc │ │ └── utils.cpython-39.pyc │ ├── auth.py │ ├── session.py │ └── utils.py ├── models │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ └── __init__.cpython-39.pyc │ ├── fastsam │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── model.cpython-311.pyc │ │ │ ├── model.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ ├── predict.cpython-39.pyc │ │ │ ├── prompt.cpython-311.pyc │ │ │ ├── prompt.cpython-39.pyc │ │ │ ├── utils.cpython-311.pyc │ │ │ ├── utils.cpython-39.pyc │ │ │ ├── val.cpython-311.pyc │ │ │ └── val.cpython-39.pyc │ │ ├── model.py │ │ ├── predict.py │ │ ├── prompt.py │ │ ├── utils.py │ │ └── val.py │ ├── nas │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── model.cpython-311.pyc │ │ │ ├── model.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ ├── predict.cpython-39.pyc │ │ │ ├── val.cpython-311.pyc │ │ │ └── val.cpython-39.pyc │ │ ├── model.py │ │ ├── predict.py │ │ └── val.py │ ├── rtdetr │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── model.cpython-311.pyc │ │ │ ├── model.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ ├── predict.cpython-39.pyc │ │ │ ├── train.cpython-311.pyc │ │ │ ├── train.cpython-39.pyc │ │ │ ├── val.cpython-311.pyc │ │ │ └── val.cpython-39.pyc │ │ ├── model.py │ │ ├── predict.py │ │ ├── train.py │ │ └── val.py │ ├── sam │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── amg.cpython-311.pyc │ │ │ ├── amg.cpython-39.pyc │ │ │ ├── build.cpython-311.pyc │ │ │ ├── build.cpython-39.pyc │ │ │ ├── model.cpython-311.pyc │ │ │ ├── model.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ └── predict.cpython-39.pyc │ │ ├── amg.py │ │ ├── build.py │ │ ├── model.py │ │ ├── modules │ │ │ ├── __pycache__ │ │ │ │ ├── __init__.cpython-39.pyc │ │ │ │ ├── decoders.cpython-311.pyc │ │ │ │ ├── decoders.cpython-39.pyc │ │ │ │ ├── encoders.cpython-311.pyc │ │ │ │ ├── encoders.cpython-39.pyc │ │ │ │ ├── sam.cpython-311.pyc │ │ │ │ ├── sam.cpython-39.pyc │ │ │ │ ├── tiny_encoder.cpython-311.pyc │ │ │ │ ├── tiny_encoder.cpython-39.pyc │ │ │ │ ├── transformer.cpython-311.pyc │ │ │ │ └── transformer.cpython-39.pyc │ │ │ ├── decoders.py │ │ │ ├── encoders.py │ │ │ ├── sam.py │ │ │ ├── tiny_encoder.py │ │ │ └── transformer.py │ │ └── predict.py │ ├── utils │ │ ├── __pycache__ │ │ │ ├── loss.cpython-311.pyc │ │ │ └── ops.cpython-311.pyc │ │ ├── loss.py │ │ └── ops.py │ └── yolo │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── model.cpython-311.pyc │ │ └── model.cpython-39.pyc │ │ ├── classify │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ ├── predict.cpython-39.pyc │ │ │ ├── train.cpython-311.pyc │ │ │ ├── train.cpython-39.pyc │ │ │ ├── val.cpython-311.pyc │ │ │ └── val.cpython-39.pyc │ │ ├── predict.py │ │ ├── train.py │ │ └── val.py │ │ ├── detect │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ ├── predict.cpython-39.pyc │ │ │ ├── train.cpython-311.pyc │ │ │ ├── train.cpython-39.pyc │ │ │ ├── val.cpython-311.pyc │ │ │ └── val.cpython-39.pyc │ │ ├── predict.py │ │ ├── train.py │ │ └── val.py │ │ ├── model.py │ │ ├── pose │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── predict.cpython-311.pyc │ │ │ ├── predict.cpython-39.pyc │ │ │ ├── train.cpython-311.pyc │ │ │ ├── train.cpython-39.pyc │ │ │ ├── val.cpython-311.pyc │ │ │ └── val.cpython-39.pyc │ │ ├── predict.py │ │ ├── train.py │ │ └── val.py │ │ └── segment │ │ ├── __init__.py │ │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── predict.cpython-311.pyc │ │ ├── predict.cpython-39.pyc │ │ ├── train.cpython-311.pyc │ │ ├── train.cpython-39.pyc │ │ ├── val.cpython-311.pyc │ │ └── val.cpython-39.pyc │ │ ├── predict.py │ │ ├── train.py │ │ └── val.py ├── nn │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── autobackend.cpython-311.pyc │ │ ├── autobackend.cpython-39.pyc │ │ ├── tasks.cpython-311.pyc │ │ └── tasks.cpython-39.pyc │ ├── autobackend.py │ ├── modules │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── block.cpython-311.pyc │ │ │ ├── block.cpython-39.pyc │ │ │ ├── conv.cpython-311.pyc │ │ │ ├── conv.cpython-39.pyc │ │ │ ├── head.cpython-311.pyc │ │ │ ├── head.cpython-39.pyc │ │ │ ├── transformer.cpython-311.pyc │ │ │ ├── transformer.cpython-39.pyc │ │ │ ├── utils.cpython-311.pyc │ │ │ └── utils.cpython-39.pyc │ │ ├── block.py │ │ ├── conv.py │ │ ├── head.py │ │ ├── transformer.py │ │ └── utils.py │ └── tasks.py ├── pipeline.png ├── trackers │ ├── README.md │ ├── __init__.py │ ├── basetrack.py │ ├── bot_sort.py │ ├── byte_tracker.py │ ├── track.py │ └── utils │ │ ├── __init__.py │ │ ├── gmc.py │ │ ├── kalman_filter.py │ │ └── matching.py └── utils │ ├── Network_functions.py │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-311.pyc │ ├── __init__.cpython-39.pyc │ ├── autobatch.cpython-311.pyc │ ├── autobatch.cpython-39.pyc │ ├── checks.cpython-311.pyc │ ├── checks.cpython-39.pyc │ ├── dist.cpython-311.pyc │ ├── dist.cpython-39.pyc │ ├── downloads.cpython-311.pyc │ ├── downloads.cpython-39.pyc │ ├── files.cpython-311.pyc │ ├── files.cpython-39.pyc │ ├── instance.cpython-311.pyc │ ├── instance.cpython-39.pyc │ ├── loss.cpython-311.pyc │ ├── loss.cpython-39.pyc │ ├── metrics.cpython-311.pyc │ ├── metrics.cpython-39.pyc │ ├── ops.cpython-311.pyc │ ├── ops.cpython-39.pyc │ ├── patches.cpython-311.pyc │ ├── patches.cpython-39.pyc │ ├── plotting.cpython-311.pyc │ ├── plotting.cpython-39.pyc │ ├── tal.cpython-311.pyc │ ├── tal.cpython-39.pyc │ ├── torch_utils.cpython-311.pyc │ └── torch_utils.cpython-39.pyc │ ├── autobatch.py │ ├── benchmarks.py │ ├── callbacks │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── base.cpython-311.pyc │ │ ├── base.cpython-39.pyc │ │ ├── clearml.cpython-311.pyc │ │ ├── clearml.cpython-39.pyc │ │ ├── comet.cpython-311.pyc │ │ ├── comet.cpython-39.pyc │ │ ├── dvc.cpython-311.pyc │ │ ├── dvc.cpython-39.pyc │ │ ├── hub.cpython-311.pyc │ │ ├── hub.cpython-39.pyc │ │ ├── mlflow.cpython-311.pyc │ │ ├── mlflow.cpython-39.pyc │ │ ├── neptune.cpython-311.pyc │ │ ├── neptune.cpython-39.pyc │ │ ├── raytune.cpython-311.pyc │ │ ├── raytune.cpython-39.pyc │ │ ├── tensorboard.cpython-311.pyc │ │ ├── tensorboard.cpython-39.pyc │ │ ├── wb.cpython-311.pyc │ │ └── wb.cpython-39.pyc │ ├── base.py │ ├── clearml.py │ ├── comet.py │ ├── dvc.py │ ├── hub.py │ ├── mlflow.py │ ├── neptune.py │ ├── raytune.py │ ├── tensorboard.py │ └── wb.py │ ├── checks.py │ ├── dist.py │ ├── downloads.py │ ├── errors.py │ ├── files.py │ ├── instance.py │ ├── loss.py │ ├── metrics.py │ ├── ops.py │ ├── patches.py │ ├── plotting.py │ ├── tal.py │ ├── torch_utils.py │ ├── triton.py │ └── tuner.py └── yoloV8N.slurm /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Y. Zambre and E. Rajkitkul and A. Mohan and J. Peeples. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Utils/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/.DS_Store -------------------------------------------------------------------------------- /Utils/Initailize_model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Functions to train/test model 4 | """ 5 | ## Python standard libraries 6 | from __future__ import print_function 7 | from __future__ import division 8 | import numpy as np 9 | import time 10 | import copy 11 | from ultralytics import YOLO, RTDETR 12 | ## PyTorch dependencies 13 | import torch 14 | import torch.nn as nn 15 | from torchvision import models 16 | import pdb 17 | 18 | def initialize_model(model_name, num_classes,feature_extract=False, 19 | use_pretrained=False): 20 | 21 | # Initialize these variables which will be set in this if statement. Each of these 22 | # variables is model specific. 23 | model_ft = None 24 | input_size = 0 25 | 26 | #Select backbone architecture 27 | if model_name == "YOLOV8N": 28 | model_ft = YOLO('yolov8n.yaml') 29 | input_size = 512 30 | 31 | elif model_name == "YOLOV8M": 32 | 33 | 34 | model_ft = YOLO('yolov8m.yaml') 35 | input_size = 512 36 | 37 | 38 | elif model_name == "YOLOV8L": 39 | 40 | 41 | model_ft = YOLO('yolov8l.yaml') 42 | input_size = 512 43 | 44 | elif model_name == "rtdetr": 45 | 46 | 47 | model_ft = RTDETR('/scratch/user/yashzambre/ultralytics/ultralytics/cfg/models/rt-detr/rtdetr-l.yaml') 48 | input_size = 512 49 | 50 | elif model_name == "convnext": 51 | model_ft = models.convnext_base(pretrained=use_pretrained) 52 | set_parameter_requires_grad(model_ft, feature_extract) 53 | num_ftrs = model_ft.classifier[2].in_features 54 | model_ft.classifier[2] = nn.Linear(model_ft.classifier[2].in_features, num_classes) 55 | input_size = 224 56 | 57 | else: 58 | raise RuntimeError('{} not implemented'.format(model_name)) 59 | 60 | 61 | return model_ft, input_size 62 | 63 | -------------------------------------------------------------------------------- /Utils/__pycache__/Compute_FDR.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Compute_FDR.cpython-39.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/Confusion_mats.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Confusion_mats.cpython-39.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/Generate_Learning_Curves.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Generate_Learning_Curves.cpython-39.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/Generate_TSNE_visual.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Generate_TSNE_visual.cpython-39.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/Network_functions.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Network_functions.cpython-311.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/Network_functions.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Network_functions.cpython-39.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/Save_Results.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/Save_Results.cpython-39.pyc -------------------------------------------------------------------------------- /Utils/__pycache__/pytorchtools.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/Utils/__pycache__/pytorchtools.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | __version__ = '8.0.120' 4 | 5 | from ultralytics.models import RTDETR, SAM, YOLO 6 | from ultralytics.models.fastsam import FastSAM 7 | from ultralytics.models.nas import NAS 8 | from ultralytics.utils import SETTINGS as settings 9 | from ultralytics.utils.checks import check_yolo as checks 10 | from ultralytics.utils.downloads import download 11 | 12 | __all__ = '__version__', 'YOLO', 'NAS', 'SAM', 'FastSAM', 'RTDETR', 'checks', 'download', 'settings' 13 | -------------------------------------------------------------------------------- /ultralytics/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/architecture.png -------------------------------------------------------------------------------- /ultralytics/assets/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/assets/bus.jpg -------------------------------------------------------------------------------- /ultralytics/assets/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/assets/zidane.jpg -------------------------------------------------------------------------------- /ultralytics/cfg/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/cfg/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/cfg/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/cfg/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/Argoverse.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI 3 | # Example usage: yolo train data=Argoverse.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── Argoverse ← downloads here (31.5 GB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/Argoverse # dataset root dir 12 | train: Argoverse-1.1/images/train/ # train images (relative to 'path') 39384 images 13 | val: Argoverse-1.1/images/val/ # val images (relative to 'path') 15062 images 14 | test: Argoverse-1.1/images/test/ # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview 15 | 16 | # Classes 17 | names: 18 | 0: person 19 | 1: bicycle 20 | 2: car 21 | 3: motorcycle 22 | 4: bus 23 | 5: truck 24 | 6: traffic_light 25 | 7: stop_sign 26 | 27 | 28 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 29 | download: | 30 | import json 31 | from tqdm import tqdm 32 | from ultralytics.utils.downloads import download 33 | from pathlib import Path 34 | 35 | def argoverse2yolo(set): 36 | labels = {} 37 | a = json.load(open(set, "rb")) 38 | for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv5 format..."): 39 | img_id = annot['image_id'] 40 | img_name = a['images'][img_id]['name'] 41 | img_label_name = f'{img_name[:-3]}txt' 42 | 43 | cls = annot['category_id'] # instance class id 44 | x_center, y_center, width, height = annot['bbox'] 45 | x_center = (x_center + width / 2) / 1920.0 # offset and scale 46 | y_center = (y_center + height / 2) / 1200.0 # offset and scale 47 | width /= 1920.0 # scale 48 | height /= 1200.0 # scale 49 | 50 | img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']] 51 | if not img_dir.exists(): 52 | img_dir.mkdir(parents=True, exist_ok=True) 53 | 54 | k = str(img_dir / img_label_name) 55 | if k not in labels: 56 | labels[k] = [] 57 | labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n") 58 | 59 | for k in labels: 60 | with open(k, "w") as f: 61 | f.writelines(labels[k]) 62 | 63 | 64 | # Download 'https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip' (deprecated S3 link) 65 | dir = Path(yaml['path']) # dataset root dir 66 | urls = ['https://drive.google.com/file/d/1st9qW3BeIwQsnR0t8mRpvbsSWIo16ACi/view?usp=drive_link'] 67 | download(urls, dir=dir) 68 | 69 | # Convert 70 | annotations_dir = 'Argoverse-HD/annotations/' 71 | (dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images') # rename 'tracking' to 'images' 72 | for d in "train.json", "val.json": 73 | argoverse2yolo(dir / annotations_dir / d) # convert Argoverse annotations to YOLO labels 74 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/DOTAv2.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # DOTA 2.0 dataset https://captain-whu.github.io/DOTA/index.html for object detection in aerial images by Wuhan University 3 | # Example usage: yolo train model=yolov8n-obb.pt data=DOTAv2.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── dota2 ← downloads here (2GB) 8 | 9 | # 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, ..] 10 | path: ../datasets/DOTAv2 # dataset root dir 11 | train: images/train # train images (relative to 'path') 1411 images 12 | val: images/val # val images (relative to 'path') 458 images 13 | test: images/test # test images (optional) 937 images 14 | 15 | # Classes for DOTA 2.0 16 | names: 17 | 0: plane 18 | 1: ship 19 | 2: storage tank 20 | 3: baseball diamond 21 | 4: tennis court 22 | 5: basketball court 23 | 6: ground track field 24 | 7: harbor 25 | 8: bridge 26 | 9: large vehicle 27 | 10: small vehicle 28 | 11: helicopter 29 | 12: roundabout 30 | 13: soccer ball field 31 | 14: swimming pool 32 | 15: container crane 33 | 16: airport 34 | 17: helipad 35 | 36 | # Download script/URL (optional) 37 | download: https://github.com/ultralytics/yolov5/releases/download/v1.0/DOTAv2.zip 38 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/GlobalWheat2020.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # Global Wheat 2020 dataset http://www.global-wheat.com/ by University of Saskatchewan 3 | # Example usage: yolo train data=GlobalWheat2020.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── GlobalWheat2020 ← downloads here (7.0 GB) 8 | 9 | num_runs: 3 10 | nc: 1 11 | # names: ['Plant'] 12 | 13 | roboflow: 14 | workspace: lab-ply6t 15 | project: plant-phenotypes 16 | version: 7 17 | license: CC BY 4.0 18 | url: https://universe.roboflow.com/lab-ply6t/plant-phenotypes/dataset/7 19 | 20 | 21 | # 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, ..] 22 | path: /scratch/user/yashzambre/ultralytics/GlobalWheat2020 # dataset root dir 23 | 24 | train: # train images (relative to 'path') 3422 images 25 | - images/arvalis_1 26 | - images/arvalis_2 27 | - images/arvalis_3 28 | - images/ethz_1 29 | - images/rres_1 30 | - images/inrae_1 31 | - images/usask_1 32 | val: # val images (relative to 'path') 748 images (WARNING: train set contains ethz_1) 33 | - images/ethz_1 34 | test: # test images (optional) 1276 images 35 | - images/utokyo_1 36 | - images/utokyo_2 37 | - images/nau_1 38 | - images/uq_1 39 | 40 | # Classes 41 | names: 42 | 0: wh 43 | 44 | 45 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/SKU-110K.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19 by Trax Retail 3 | # Example usage: yolo train data=SKU-110K.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── SKU-110K ← downloads here (13.6 GB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/SKU-110K # dataset root dir 12 | train: train.txt # train images (relative to 'path') 8219 images 13 | val: val.txt # val images (relative to 'path') 588 images 14 | test: test.txt # test images (optional) 2936 images 15 | 16 | # Classes 17 | names: 18 | 0: object 19 | 20 | 21 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 22 | download: | 23 | import shutil 24 | from pathlib import Path 25 | 26 | import numpy as np 27 | import pandas as pd 28 | from tqdm import tqdm 29 | 30 | from ultralytics.utils.downloads import download 31 | from ultralytics.utils.ops import xyxy2xywh 32 | 33 | # Download 34 | dir = Path(yaml['path']) # dataset root dir 35 | parent = Path(dir.parent) # download dir 36 | urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz'] 37 | download(urls, dir=parent) 38 | 39 | # Rename directories 40 | if dir.exists(): 41 | shutil.rmtree(dir) 42 | (parent / 'SKU110K_fixed').rename(dir) # rename dir 43 | (dir / 'labels').mkdir(parents=True, exist_ok=True) # create labels dir 44 | 45 | # Convert labels 46 | names = 'image', 'x1', 'y1', 'x2', 'y2', 'class', 'image_width', 'image_height' # column names 47 | for d in 'annotations_train.csv', 'annotations_val.csv', 'annotations_test.csv': 48 | x = pd.read_csv(dir / 'annotations' / d, names=names).values # annotations 49 | images, unique_images = x[:, 0], np.unique(x[:, 0]) 50 | with open((dir / d).with_suffix('.txt').__str__().replace('annotations_', ''), 'w') as f: 51 | f.writelines(f'./images/{s}\n' for s in unique_images) 52 | for im in tqdm(unique_images, desc=f'Converting {dir / d}'): 53 | cls = 0 # single-class dataset 54 | with open((dir / 'labels' / im).with_suffix('.txt'), 'a') as f: 55 | for r in x[images == im]: 56 | w, h = r[6], r[7] # image width, height 57 | xywh = xyxy2xywh(np.array([[r[1] / w, r[2] / h, r[3] / w, r[4] / h]]))[0] # instance 58 | f.write(f"{cls} {xywh[0]:.5f} {xywh[1]:.5f} {xywh[2]:.5f} {xywh[3]:.5f}\n") # write label 59 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/VisDrone.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO , AGPL-3.0 license 2 | # VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University 3 | # Example usage: yolo train data=VisDrone.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── VisDrone ← downloads here (2.3 GB) 8 | 9 | num_runs: 3 10 | nc: 10 11 | # names: ['Plant'] 12 | 13 | roboflow: 14 | workspace: lab-ply6t 15 | project: plant-phenotypes 16 | version: 7 17 | license: CC BY 4.0 18 | url: https://universe.roboflow.com/lab-ply6t/plant-phenotypes/dataset/7 19 | # 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, ..] 20 | path: /scratch/user/yashzambre/ultralytics/VisDrone 21 | train: VisDrone2019-DET-train/images # train images (relative to 'path') 6471 images 22 | val: VisDrone2019-DET-val/images # val images (relative to 'path') 548 images 23 | test: VisDrone2019-DET-test-dev/images # test images (optional) 1610 images 24 | 25 | # Classes 26 | names: 27 | 0: pedestrian 28 | 1: people 29 | 2: bicycle 30 | 3: car 31 | 4: van 32 | 5: truck 33 | 6: tricycle 34 | 7: awning-tricycle 35 | 8: bus 36 | 9: motor 37 | 38 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 39 | download: | 40 | import os 41 | from pathlib import Path 42 | 43 | from ultralytics.utils.downloads import download 44 | 45 | def visdrone2yolo(dir): 46 | from PIL import Image 47 | from tqdm import tqdm 48 | 49 | def convert_box(size, box): 50 | # Convert VisDrone box to YOLO xywh box 51 | dw = 1. / size[0] 52 | dh = 1. / size[1] 53 | return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh 54 | 55 | (dir / 'labels').mkdir(parents=True, exist_ok=True) # make labels directory 56 | pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}') 57 | for f in pbar: 58 | img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size 59 | lines = [] 60 | with open(f, 'r') as file: # read annotation.txt 61 | for row in [x.split(',') for x in file.read().strip().splitlines()]: 62 | if row[4] == '0': # VisDrone 'ignored regions' class 0 63 | continue 64 | cls = int(row[5]) - 1 65 | box = convert_box(img_size, tuple(map(int, row[:4]))) 66 | lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n") 67 | with open(str(f).replace(f'{os.sep}annotations{os.sep}', f'{os.sep}labels{os.sep}'), 'w') as fl: 68 | fl.writelines(lines) # write label.txt 69 | 70 | 71 | # Download 72 | dir = Path(yaml['path']) # dataset root dir 73 | urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip', 74 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip', 75 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip', 76 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip'] 77 | download(urls, dir=dir, curl=True, threads=4) 78 | 79 | # Convert 80 | for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev': 81 | visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels 82 | 83 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco-pose.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO 2017 dataset http://cocodataset.org by Microsoft 3 | # Example usage: yolo train data=coco-pose.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco-pose ← downloads here (20.1 GB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/coco-pose # dataset root dir 12 | train: train2017.txt # train images (relative to 'path') 118287 images 13 | val: val2017.txt # val images (relative to 'path') 5000 images 14 | test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794 15 | 16 | # Keypoints 17 | kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) 18 | flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] 19 | 20 | # Classes 21 | names: 22 | 0: person 23 | 24 | # Download script/URL (optional) 25 | download: | 26 | from ultralytics.utils.downloads import download 27 | from pathlib import Path 28 | 29 | # Download labels 30 | dir = Path(yaml['path']) # dataset root dir 31 | url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/' 32 | urls = [url + 'coco2017labels-pose.zip'] # labels 33 | download(urls, dir=dir.parent) 34 | # Download data 35 | urls = ['http://images.cocodataset.org/zips/train2017.zip', # 19G, 118k images 36 | 'http://images.cocodataset.org/zips/val2017.zip', # 1G, 5k images 37 | 'http://images.cocodataset.org/zips/test2017.zip'] # 7G, 41k images (optional) 38 | download(urls, dir=dir / 'images', threads=3) 39 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO 2017 dataset http://cocodataset.org by Microsoft 3 | # Example usage: yolo train data=coco.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco ← downloads here (20.1 GB) 8 | 9 | num_runs: 3 10 | nc: 80 11 | # names: ['Plant'] 12 | 13 | roboflow: 14 | workspace: lab-ply6t 15 | project: plant-phenotypes 16 | version: 7 17 | license: CC BY 4.0 18 | url: https://universe.roboflow.com/lab-ply6t/plant-phenotypes/dataset/7 19 | # 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, ..] 20 | path: /scratch/user/yashzambre/ultralytics/coco # dataset root dir 21 | train: images/train2017 # train images (relative to 'path') 118287 images 22 | val: images/val2017 # val images (relative to 'path') 5000 images 23 | test: images/test2017 # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794 24 | 25 | # Classes 26 | names: 27 | 0: person 28 | 1: bicycle 29 | 2: car 30 | 3: motorcycle 31 | 4: airplane 32 | 5: bus 33 | 6: train 34 | 7: truck 35 | 8: boat 36 | 9: traffic light 37 | 10: fire hydrant 38 | 11: stop sign 39 | 12: parking meter 40 | 13: bench 41 | 14: bird 42 | 15: cat 43 | 16: dog 44 | 17: horse 45 | 18: sheep 46 | 19: cow 47 | 20: elephant 48 | 21: bear 49 | 22: zebra 50 | 23: giraffe 51 | 24: backpack 52 | 25: umbrella 53 | 26: handbag 54 | 27: tie 55 | 28: suitcase 56 | 29: frisbee 57 | 30: skis 58 | 31: snowboard 59 | 32: sports ball 60 | 33: kite 61 | 34: baseball bat 62 | 35: baseball glove 63 | 36: skateboard 64 | 37: surfboard 65 | 38: tennis racket 66 | 39: bottle 67 | 40: wine glass 68 | 41: cup 69 | 42: fork 70 | 43: knife 71 | 44: spoon 72 | 45: bowl 73 | 46: banana 74 | 47: apple 75 | 48: sandwich 76 | 49: orange 77 | 50: broccoli 78 | 51: carrot 79 | 52: hot dog 80 | 53: pizza 81 | 54: donut 82 | 55: cake 83 | 56: chair 84 | 57: couch 85 | 58: potted plant 86 | 59: bed 87 | 60: dining table 88 | 61: toilet 89 | 62: tv 90 | 63: laptop 91 | 64: mouse 92 | 65: remote 93 | 66: keyboard 94 | 67: cell phone 95 | 68: microwave 96 | 69: oven 97 | 70: toaster 98 | 71: sink 99 | 72: refrigerator 100 | 73: book 101 | 74: clock 102 | 75: vase 103 | 76: scissors 104 | 77: teddy bear 105 | 78: hair drier 106 | 79: toothbrush 107 | 108 | 109 | 110 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco128-seg.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics 3 | # Example usage: yolo train data=coco128.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco128-seg ← downloads here (7 MB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/coco128-seg # dataset root dir 12 | train: images/train2017 # train images (relative to 'path') 128 images 13 | val: images/train2017 # val images (relative to 'path') 128 images 14 | test: # test images (optional) 15 | 16 | # Classes 17 | names: 18 | 0: person 19 | 1: bicycle 20 | 2: car 21 | 3: motorcycle 22 | 4: airplane 23 | 5: bus 24 | 6: train 25 | 7: truck 26 | 8: boat 27 | 9: traffic light 28 | 10: fire hydrant 29 | 11: stop sign 30 | 12: parking meter 31 | 13: bench 32 | 14: bird 33 | 15: cat 34 | 16: dog 35 | 17: horse 36 | 18: sheep 37 | 19: cow 38 | 20: elephant 39 | 21: bear 40 | 22: zebra 41 | 23: giraffe 42 | 24: backpack 43 | 25: umbrella 44 | 26: handbag 45 | 27: tie 46 | 28: suitcase 47 | 29: frisbee 48 | 30: skis 49 | 31: snowboard 50 | 32: sports ball 51 | 33: kite 52 | 34: baseball bat 53 | 35: baseball glove 54 | 36: skateboard 55 | 37: surfboard 56 | 38: tennis racket 57 | 39: bottle 58 | 40: wine glass 59 | 41: cup 60 | 42: fork 61 | 43: knife 62 | 44: spoon 63 | 45: bowl 64 | 46: banana 65 | 47: apple 66 | 48: sandwich 67 | 49: orange 68 | 50: broccoli 69 | 51: carrot 70 | 52: hot dog 71 | 53: pizza 72 | 54: donut 73 | 55: cake 74 | 56: chair 75 | 57: couch 76 | 58: potted plant 77 | 59: bed 78 | 60: dining table 79 | 61: toilet 80 | 62: tv 81 | 63: laptop 82 | 64: mouse 83 | 65: remote 84 | 66: keyboard 85 | 67: cell phone 86 | 68: microwave 87 | 69: oven 88 | 70: toaster 89 | 71: sink 90 | 72: refrigerator 91 | 73: book 92 | 74: clock 93 | 75: vase 94 | 76: scissors 95 | 77: teddy bear 96 | 78: hair drier 97 | 79: toothbrush 98 | 99 | 100 | # Download script/URL (optional) 101 | download: https://ultralytics.com/assets/coco128-seg.zip 102 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco128.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics 3 | # Example usage: yolo train data=coco128.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco128 ← downloads here (7 MB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/coco128 # dataset root dir 12 | train: images/train2017 # train images (relative to 'path') 128 images 13 | val: images/train2017 # val images (relative to 'path') 128 images 14 | test: # test images (optional) 15 | 16 | # Classes 17 | names: 18 | 0: person 19 | 1: bicycle 20 | 2: car 21 | 3: motorcycle 22 | 4: airplane 23 | 5: bus 24 | 6: train 25 | 7: truck 26 | 8: boat 27 | 9: traffic light 28 | 10: fire hydrant 29 | 11: stop sign 30 | 12: parking meter 31 | 13: bench 32 | 14: bird 33 | 15: cat 34 | 16: dog 35 | 17: horse 36 | 18: sheep 37 | 19: cow 38 | 20: elephant 39 | 21: bear 40 | 22: zebra 41 | 23: giraffe 42 | 24: backpack 43 | 25: umbrella 44 | 26: handbag 45 | 27: tie 46 | 28: suitcase 47 | 29: frisbee 48 | 30: skis 49 | 31: snowboard 50 | 32: sports ball 51 | 33: kite 52 | 34: baseball bat 53 | 35: baseball glove 54 | 36: skateboard 55 | 37: surfboard 56 | 38: tennis racket 57 | 39: bottle 58 | 40: wine glass 59 | 41: cup 60 | 42: fork 61 | 43: knife 62 | 44: spoon 63 | 45: bowl 64 | 46: banana 65 | 47: apple 66 | 48: sandwich 67 | 49: orange 68 | 50: broccoli 69 | 51: carrot 70 | 52: hot dog 71 | 53: pizza 72 | 54: donut 73 | 55: cake 74 | 56: chair 75 | 57: couch 76 | 58: potted plant 77 | 59: bed 78 | 60: dining table 79 | 61: toilet 80 | 62: tv 81 | 63: laptop 82 | 64: mouse 83 | 65: remote 84 | 66: keyboard 85 | 67: cell phone 86 | 68: microwave 87 | 69: oven 88 | 70: toaster 89 | 71: sink 90 | 72: refrigerator 91 | 73: book 92 | 74: clock 93 | 75: vase 94 | 76: scissors 95 | 77: teddy bear 96 | 78: hair drier 97 | 79: toothbrush 98 | 99 | 100 | # Download script/URL (optional) 101 | download: https://ultralytics.com/assets/coco128.zip 102 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco8-pose.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO8-pose dataset (first 8 images from COCO train2017) by Ultralytics 3 | # Example usage: yolo train data=coco8-pose.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco8-pose ← downloads here (1 MB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/coco8-pose # dataset root dir 12 | train: images/train # train images (relative to 'path') 4 images 13 | val: images/val # val images (relative to 'path') 4 images 14 | test: # test images (optional) 15 | 16 | # Keypoints 17 | kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) 18 | flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15] 19 | 20 | # Classes 21 | names: 22 | 0: person 23 | 24 | # Download script/URL (optional) 25 | download: https://ultralytics.com/assets/coco8-pose.zip 26 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco8-seg.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO8-seg dataset (first 8 images from COCO train2017) by Ultralytics 3 | # Example usage: yolo train data=coco8-seg.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco8-seg ← downloads here (1 MB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/coco8-seg # dataset root dir 12 | train: images/train # train images (relative to 'path') 4 images 13 | val: images/val # val images (relative to 'path') 4 images 14 | test: # test images (optional) 15 | 16 | # Classes 17 | names: 18 | 0: person 19 | 1: bicycle 20 | 2: car 21 | 3: motorcycle 22 | 4: airplane 23 | 5: bus 24 | 6: train 25 | 7: truck 26 | 8: boat 27 | 9: traffic light 28 | 10: fire hydrant 29 | 11: stop sign 30 | 12: parking meter 31 | 13: bench 32 | 14: bird 33 | 15: cat 34 | 16: dog 35 | 17: horse 36 | 18: sheep 37 | 19: cow 38 | 20: elephant 39 | 21: bear 40 | 22: zebra 41 | 23: giraffe 42 | 24: backpack 43 | 25: umbrella 44 | 26: handbag 45 | 27: tie 46 | 28: suitcase 47 | 29: frisbee 48 | 30: skis 49 | 31: snowboard 50 | 32: sports ball 51 | 33: kite 52 | 34: baseball bat 53 | 35: baseball glove 54 | 36: skateboard 55 | 37: surfboard 56 | 38: tennis racket 57 | 39: bottle 58 | 40: wine glass 59 | 41: cup 60 | 42: fork 61 | 43: knife 62 | 44: spoon 63 | 45: bowl 64 | 46: banana 65 | 47: apple 66 | 48: sandwich 67 | 49: orange 68 | 50: broccoli 69 | 51: carrot 70 | 52: hot dog 71 | 53: pizza 72 | 54: donut 73 | 55: cake 74 | 56: chair 75 | 57: couch 76 | 58: potted plant 77 | 59: bed 78 | 60: dining table 79 | 61: toilet 80 | 62: tv 81 | 63: laptop 82 | 64: mouse 83 | 65: remote 84 | 66: keyboard 85 | 67: cell phone 86 | 68: microwave 87 | 69: oven 88 | 70: toaster 89 | 71: sink 90 | 72: refrigerator 91 | 73: book 92 | 74: clock 93 | 75: vase 94 | 76: scissors 95 | 77: teddy bear 96 | 78: hair drier 97 | 79: toothbrush 98 | 99 | 100 | # Download script/URL (optional) 101 | download: https://ultralytics.com/assets/coco8-seg.zip 102 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/coco8.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # COCO8 dataset (first 8 images from COCO train2017) by Ultralytics 3 | # Example usage: yolo train data=coco8.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── coco8 ← downloads here (1 MB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/coco8 # dataset root dir 12 | train: images/train # train images (relative to 'path') 4 images 13 | val: images/val # val images (relative to 'path') 4 images 14 | test: # test images (optional) 15 | 16 | # Classes 17 | names: 18 | 0: person 19 | 1: bicycle 20 | 2: car 21 | 3: motorcycle 22 | 4: airplane 23 | 5: bus 24 | 6: train 25 | 7: truck 26 | 8: boat 27 | 9: traffic light 28 | 10: fire hydrant 29 | 11: stop sign 30 | 12: parking meter 31 | 13: bench 32 | 14: bird 33 | 15: cat 34 | 16: dog 35 | 17: horse 36 | 18: sheep 37 | 19: cow 38 | 20: elephant 39 | 21: bear 40 | 22: zebra 41 | 23: giraffe 42 | 24: backpack 43 | 25: umbrella 44 | 26: handbag 45 | 27: tie 46 | 28: suitcase 47 | 29: frisbee 48 | 30: skis 49 | 31: snowboard 50 | 32: sports ball 51 | 33: kite 52 | 34: baseball bat 53 | 35: baseball glove 54 | 36: skateboard 55 | 37: surfboard 56 | 38: tennis racket 57 | 39: bottle 58 | 40: wine glass 59 | 41: cup 60 | 42: fork 61 | 43: knife 62 | 44: spoon 63 | 45: bowl 64 | 46: banana 65 | 47: apple 66 | 48: sandwich 67 | 49: orange 68 | 50: broccoli 69 | 51: carrot 70 | 52: hot dog 71 | 53: pizza 72 | 54: donut 73 | 55: cake 74 | 56: chair 75 | 57: couch 76 | 58: potted plant 77 | 59: bed 78 | 60: dining table 79 | 61: toilet 80 | 62: tv 81 | 63: laptop 82 | 64: mouse 83 | 65: remote 84 | 66: keyboard 85 | 67: cell phone 86 | 68: microwave 87 | 69: oven 88 | 70: toaster 89 | 71: sink 90 | 72: refrigerator 91 | 73: book 92 | 74: clock 93 | 75: vase 94 | 76: scissors 95 | 77: teddy bear 96 | 78: hair drier 97 | 79: toothbrush 98 | 99 | 100 | # Download script/URL (optional) 101 | download: https://ultralytics.com/assets/coco8.zip 102 | -------------------------------------------------------------------------------- /ultralytics/cfg/datasets/tiger-pose.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # Tiger Pose dataset by Ultralytics 3 | # Example usage: yolo train data=tiger-pose.yaml 4 | # parent 5 | # ├── ultralytics 6 | # └── datasets 7 | # └── tiger-pose ← downloads here (75.3 MB) 8 | 9 | 10 | # 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, ..] 11 | path: ../datasets/tiger-pose # dataset root dir 12 | train: train # train images (relative to 'path') 210 images 13 | val: val # val images (relative to 'path') 53 images 14 | 15 | # Keypoints 16 | kpt_shape: [12, 2] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) 17 | flip_idx: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] 18 | 19 | # Classes 20 | names: 21 | 0: tiger 22 | 23 | # Download script/URL (optional) 24 | download: https://ultralytics.com/assets/tiger-pose.zip 25 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/README.md: -------------------------------------------------------------------------------- 1 | ## Models 2 | 3 | Welcome to the Ultralytics Models directory! Here you will find a wide variety of pre-configured model configuration files (`*.yaml`s) that can be used to create custom YOLO models. The models in this directory have been expertly crafted and fine-tuned by the Ultralytics team to provide the best performance for a wide range of object detection and image segmentation tasks. 4 | 5 | These model configurations cover a wide range of scenarios, from simple object detection to more complex tasks like instance segmentation and object tracking. They are also designed to run efficiently on a variety of hardware platforms, from CPUs to GPUs. Whether you are a seasoned machine learning practitioner or just getting started with YOLO, this directory provides a great starting point for your custom model development needs. 6 | 7 | To get started, simply browse through the models in this directory and find one that best suits your needs. Once you've selected a model, you can use the provided `*.yaml` file to train and deploy your custom YOLO model with ease. See full details at the Ultralytics [Docs](https://docs.ultralytics.com/models), and if you need help or have any questions, feel free to reach out to the Ultralytics team for support. So, don't wait, start creating your custom YOLO model now! 8 | 9 | ### Usage 10 | 11 | Model `*.yaml` files may be used directly in the Command Line Interface (CLI) with a `yolo` command: 12 | 13 | ```bash 14 | yolo task=detect mode=train model=yolov8n.yaml data=coco128.yaml epochs=100 15 | ``` 16 | 17 | They may also be used directly in a Python environment, and accepts the same 18 | [arguments](https://docs.ultralytics.com/usage/cfg/) as in the CLI example above: 19 | 20 | ```python 21 | from ultralytics import YOLO 22 | 23 | model = YOLO("model.yaml") # build a YOLOv8n model from scratch 24 | # YOLO("model.pt") use pre-trained model if available 25 | model.info() # display model information 26 | model.train(data="coco128.yaml", epochs=100) # train the model 27 | ``` 28 | 29 | ## Pre-trained Model Architectures 30 | 31 | Ultralytics supports many model architectures. Visit https://docs.ultralytics.com/models to view detailed information and usage. Any of these models can be used by loading their configs or pretrained checkpoints if available. 32 | 33 | ## Contribute New Models 34 | 35 | Have you trained a new YOLO variant or achieved state-of-the-art performance with specific tuning? We'd love to showcase your work in our Models section! Contributions from the community in the form of new models, architectures, or optimizations are highly valued and can significantly enrich our repository. 36 | 37 | By contributing to this section, you're helping us offer a wider array of model choices and configurations to the community. It's a fantastic way to share your knowledge and expertise while making the Ultralytics YOLO ecosystem even more versatile. 38 | 39 | To get started, please consult our [Contributing Guide](https://docs.ultralytics.com/help/contributing) for step-by-step instructions on how to submit a Pull Request (PR) 🛠️. Your contributions are eagerly awaited! 40 | 41 | Let's join hands to extend the range and capabilities of the Ultralytics YOLO models 🙏! 42 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/rt-detr/rtdetr-l.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # RT-DETR-l object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/rtdetr 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | l: [1.00, 1.00, 1024] 9 | 10 | backbone: 11 | # [from, repeats, module, args] 12 | - [-1,1,SpatialTransformer ,[3]] 13 | - [-1, 1, HGStem, [32, 48]] # 0-P2/4 14 | - [-1, 6, HGBlock, [48, 128, 3]] # stage 1 15 | 16 | - [-1, 1, DWConv, [128, 3, 2, 1, False]] # 2-P3/8 17 | - [-1, 6, HGBlock, [96, 512, 3]] # stage 2 18 | 19 | - [-1, 1, DWConv, [512, 3, 2, 1, False]] # 4-P3/16 20 | - [-1, 6, HGBlock, [192, 1024, 5, True, False]] # cm, c2, k, light, shortcut 21 | - [-1, 6, HGBlock, [192, 1024, 5, True, True]] 22 | - [-1, 6, HGBlock, [192, 1024, 5, True, True]] # stage 3 23 | 24 | - [-1, 1, DWConv, [1024, 3, 2, 1, False]] # 8-P4/32 25 | - [-1, 6, HGBlock, [384, 2048, 5, True, False]] # stage 4 26 | 27 | head: 28 | - [-1, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 10 input_proj.2 29 | - [-1, 1, AIFI, [1024, 8]] 30 | - [-1, 1, Conv, [256, 1, 1]] # 12, Y5, lateral_convs.0 31 | 32 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 33 | - [7, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 14 input_proj.1 34 | - [[-2, -1], 1, Concat, [1]] 35 | - [-1, 3, RepC3, [256]] # 16, fpn_blocks.0 36 | - [-1, 1, Conv, [256, 1, 1]] # 17, Y4, lateral_convs.1 37 | 38 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 39 | - [3, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 19 input_proj.0 40 | - [[-2, -1], 1, Concat, [1]] # cat backbone P4 41 | - [-1, 3, RepC3, [256]] # X3 (21), fpn_blocks.1 42 | 43 | - [-1, 1, Conv, [256, 3, 2]] # 22, downsample_convs.0 44 | - [[-1, 17], 1, Concat, [1]] # cat Y4 45 | - [-1, 3, RepC3, [256]] # F4 (24), pan_blocks.0 46 | 47 | - [-1, 1, Conv, [256, 3, 2]] # 25, downsample_convs.1 48 | - [[-1, 12], 1, Concat, [1]] # cat Y5 49 | - [-1, 3, RepC3, [256]] # F5 (27), pan_blocks.1 50 | 51 | - [[21, 24, 27], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5) 52 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/rt-detr/rtdetr-x.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # RT-DETR-x object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/rtdetr 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | x: [1.00, 1.00, 2048] 9 | 10 | backbone: 11 | # [from, repeats, module, args] 12 | - [-1, 1, HGStem, [32, 64]] # 0-P2/4 13 | - [-1, 6, HGBlock, [64, 128, 3]] # stage 1 14 | 15 | - [-1, 1, DWConv, [128, 3, 2, 1, False]] # 2-P3/8 16 | - [-1, 6, HGBlock, [128, 512, 3]] 17 | - [-1, 6, HGBlock, [128, 512, 3, False, True]] # 4-stage 2 18 | 19 | - [-1, 1, DWConv, [512, 3, 2, 1, False]] # 5-P3/16 20 | - [-1, 6, HGBlock, [256, 1024, 5, True, False]] # cm, c2, k, light, shortcut 21 | - [-1, 6, HGBlock, [256, 1024, 5, True, True]] 22 | - [-1, 6, HGBlock, [256, 1024, 5, True, True]] 23 | - [-1, 6, HGBlock, [256, 1024, 5, True, True]] 24 | - [-1, 6, HGBlock, [256, 1024, 5, True, True]] # 10-stage 3 25 | 26 | - [-1, 1, DWConv, [1024, 3, 2, 1, False]] # 11-P4/32 27 | - [-1, 6, HGBlock, [512, 2048, 5, True, False]] 28 | - [-1, 6, HGBlock, [512, 2048, 5, True, True]] # 13-stage 4 29 | 30 | head: 31 | - [-1, 1, Conv, [384, 1, 1, None, 1, 1, False]] # 14 input_proj.2 32 | - [-1, 1, AIFI, [2048, 8]] 33 | - [-1, 1, Conv, [384, 1, 1]] # 16, Y5, lateral_convs.0 34 | 35 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 36 | - [10, 1, Conv, [384, 1, 1, None, 1, 1, False]] # 18 input_proj.1 37 | - [[-2, -1], 1, Concat, [1]] 38 | - [-1, 3, RepC3, [384]] # 20, fpn_blocks.0 39 | - [-1, 1, Conv, [384, 1, 1]] # 21, Y4, lateral_convs.1 40 | 41 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 42 | - [4, 1, Conv, [384, 1, 1, None, 1, 1, False]] # 23 input_proj.0 43 | - [[-2, -1], 1, Concat, [1]] # cat backbone P4 44 | - [-1, 3, RepC3, [384]] # X3 (25), fpn_blocks.1 45 | 46 | - [-1, 1, Conv, [384, 3, 2]] # 26, downsample_convs.0 47 | - [[-1, 21], 1, Concat, [1]] # cat Y4 48 | - [-1, 3, RepC3, [384]] # F4 (28), pan_blocks.0 49 | 50 | - [-1, 1, Conv, [384, 3, 2]] # 29, downsample_convs.1 51 | - [[-1, 16], 1, Concat, [1]] # cat Y5 52 | - [-1, 3, RepC3, [384]] # F5 (31), pan_blocks.1 53 | 54 | - [[25, 28, 31], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5) 55 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v3/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv3-SPP object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/yolov3 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | depth_multiple: 1.0 # model depth multiple 7 | width_multiple: 1.0 # layer channel multiple 8 | 9 | # darknet53 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [32, 3, 1]], # 0 13 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 14 | [-1, 1, Bottleneck, [64]], 15 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 16 | [-1, 2, Bottleneck, [128]], 17 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 18 | [-1, 8, Bottleneck, [256]], 19 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 20 | [-1, 8, Bottleneck, [512]], 21 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 22 | [-1, 4, Bottleneck, [1024]], # 10 23 | ] 24 | 25 | # YOLOv3-SPP head 26 | head: 27 | [[-1, 1, Bottleneck, [1024, False]], 28 | [-1, 1, SPP, [512, [5, 9, 13]]], 29 | [-1, 1, Conv, [1024, 3, 1]], 30 | [-1, 1, Conv, [512, 1, 1]], 31 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 32 | 33 | [-2, 1, Conv, [256, 1, 1]], 34 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 35 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 36 | [-1, 1, Bottleneck, [512, False]], 37 | [-1, 1, Bottleneck, [512, False]], 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 40 | 41 | [-2, 1, Conv, [128, 1, 1]], 42 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 43 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 44 | [-1, 1, Bottleneck, [256, False]], 45 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 46 | 47 | [[27, 22, 15], 1, Detect, [nc]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v3/yolov3-tiny.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv3-tiny object detection model with P4-P5 outputs. For details see https://docs.ultralytics.com/models/yolov3 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | depth_multiple: 1.0 # model depth multiple 7 | width_multiple: 1.0 # layer channel multiple 8 | 9 | # YOLOv3-tiny backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [16, 3, 1]], # 0 13 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 1-P1/2 14 | [-1, 1, Conv, [32, 3, 1]], 15 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 3-P2/4 16 | [-1, 1, Conv, [64, 3, 1]], 17 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 5-P3/8 18 | [-1, 1, Conv, [128, 3, 1]], 19 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 7-P4/16 20 | [-1, 1, Conv, [256, 3, 1]], 21 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 9-P5/32 22 | [-1, 1, Conv, [512, 3, 1]], 23 | [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]], # 11 24 | [-1, 1, nn.MaxPool2d, [2, 1, 0]], # 12 25 | ] 26 | 27 | # YOLOv3-tiny head 28 | head: 29 | [[-1, 1, Conv, [1024, 3, 1]], 30 | [-1, 1, Conv, [256, 1, 1]], 31 | [-1, 1, Conv, [512, 3, 1]], # 15 (P5/32-large) 32 | 33 | [-2, 1, Conv, [128, 1, 1]], 34 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 35 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 36 | [-1, 1, Conv, [256, 3, 1]], # 19 (P4/16-medium) 37 | 38 | [[19, 15], 1, Detect, [nc]], # Detect(P4, P5) 39 | ] 40 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v3/yolov3.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv3 object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/yolov3 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | depth_multiple: 1.0 # model depth multiple 7 | width_multiple: 1.0 # layer channel multiple 8 | 9 | # darknet53 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [32, 3, 1]], # 0 13 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 14 | [-1, 1, Bottleneck, [64]], 15 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 16 | [-1, 2, Bottleneck, [128]], 17 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 18 | [-1, 8, Bottleneck, [256]], 19 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 20 | [-1, 8, Bottleneck, [512]], 21 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 22 | [-1, 4, Bottleneck, [1024]], # 10 23 | ] 24 | 25 | # YOLOv3 head 26 | head: 27 | [[-1, 1, Bottleneck, [1024, False]], 28 | [-1, 1, Conv, [512, 1, 1]], 29 | [-1, 1, Conv, [1024, 3, 1]], 30 | [-1, 1, Conv, [512, 1, 1]], 31 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 32 | 33 | [-2, 1, Conv, [256, 1, 1]], 34 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 35 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 36 | [-1, 1, Bottleneck, [512, False]], 37 | [-1, 1, Bottleneck, [512, False]], 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 40 | 41 | [-2, 1, Conv, [128, 1, 1]], 42 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 43 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 44 | [-1, 1, Bottleneck, [256, False]], 45 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 46 | 47 | [[27, 22, 15], 1, Detect, [nc]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v5/yolov5-p6.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv5 object detection model with P3-P6 outputs. For details see https://docs.ultralytics.com/models/yolov5 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov5n-p6.yaml' will call yolov5-p6.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 1024] 11 | l: [1.00, 1.00, 1024] 12 | x: [1.33, 1.25, 1024] 13 | 14 | # YOLOv5 v6.0 backbone 15 | backbone: 16 | # [from, number, module, args] 17 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 18 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 19 | [-1, 3, C3, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 21 | [-1, 6, C3, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 23 | [-1, 9, C3, [512]], 24 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 25 | [-1, 3, C3, [768]], 26 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 27 | [-1, 3, C3, [1024]], 28 | [-1, 1, SPPF, [1024, 5]], # 11 29 | ] 30 | 31 | # YOLOv5 v6.0 head 32 | head: 33 | [[-1, 1, Conv, [768, 1, 1]], 34 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 35 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 36 | [-1, 3, C3, [768, False]], # 15 37 | 38 | [-1, 1, Conv, [512, 1, 1]], 39 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 40 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 41 | [-1, 3, C3, [512, False]], # 19 42 | 43 | [-1, 1, Conv, [256, 1, 1]], 44 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 45 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 46 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 47 | 48 | [-1, 1, Conv, [256, 3, 2]], 49 | [[-1, 20], 1, Concat, [1]], # cat head P4 50 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 51 | 52 | [-1, 1, Conv, [512, 3, 2]], 53 | [[-1, 16], 1, Concat, [1]], # cat head P5 54 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 55 | 56 | [-1, 1, Conv, [768, 3, 2]], 57 | [[-1, 12], 1, Concat, [1]], # cat head P6 58 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 59 | 60 | [[23, 26, 29, 32], 1, Detect, [nc]], # Detect(P3, P4, P5, P6) 61 | ] 62 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v5/yolov5.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv5 object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/yolov5 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov5n.yaml' will call yolov5.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 1024] 11 | l: [1.00, 1.00, 1024] 12 | x: [1.33, 1.25, 1024] 13 | 14 | # YOLOv5 v6.0 backbone 15 | backbone: 16 | # [from, number, module, args] 17 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 18 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 19 | [-1, 3, C3, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 21 | [-1, 6, C3, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 23 | [-1, 9, C3, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 25 | [-1, 3, C3, [1024]], 26 | [-1, 1, SPPF, [1024, 5]], # 9 27 | ] 28 | 29 | # YOLOv5 v6.0 head 30 | head: 31 | [[-1, 1, Conv, [512, 1, 1]], 32 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 33 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 34 | [-1, 3, C3, [512, False]], # 13 35 | 36 | [-1, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 39 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 40 | 41 | [-1, 1, Conv, [256, 3, 2]], 42 | [[-1, 14], 1, Concat, [1]], # cat head P4 43 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 44 | 45 | [-1, 1, Conv, [512, 3, 2]], 46 | [[-1, 10], 1, Concat, [1]], # cat head P5 47 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 48 | 49 | [[17, 20, 23], 1, Detect, [nc]], # Detect(P3, P4, P5) 50 | ] 51 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v6/yolov6.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv6 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/models/yolov6 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | activation: nn.ReLU() # (optional) model default activation function 7 | scales: # model compound scaling constants, i.e. 'model=yolov6n.yaml' will call yolov8.yaml with scale 'n' 8 | # [depth, width, max_channels] 9 | n: [0.33, 0.25, 1024] 10 | s: [0.33, 0.50, 1024] 11 | m: [0.67, 0.75, 768] 12 | l: [1.00, 1.00, 512] 13 | x: [1.00, 1.25, 512] 14 | 15 | # YOLOv6-3.0s backbone 16 | backbone: 17 | # [from, repeats, module, args] 18 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 19 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 20 | - [-1, 6, Conv, [128, 3, 1]] 21 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 22 | - [-1, 12, Conv, [256, 3, 1]] 23 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 24 | - [-1, 18, Conv, [512, 3, 1]] 25 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 26 | - [-1, 6, Conv, [1024, 3, 1]] 27 | - [-1, 1, SPPF, [1024, 5]] # 9 28 | 29 | # YOLOv6-3.0s head 30 | head: 31 | - [-1, 1, Conv, [256, 1, 1]] 32 | - [-1, 1, nn.ConvTranspose2d, [256, 2, 2, 0]] 33 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 34 | - [-1, 1, Conv, [256, 3, 1]] 35 | - [-1, 9, Conv, [256, 3, 1]] # 14 36 | 37 | - [-1, 1, Conv, [128, 1, 1]] 38 | - [-1, 1, nn.ConvTranspose2d, [128, 2, 2, 0]] 39 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 40 | - [-1, 1, Conv, [128, 3, 1]] 41 | - [-1, 9, Conv, [128, 3, 1]] # 19 42 | 43 | - [-1, 1, Conv, [128, 3, 2]] 44 | - [[-1, 15], 1, Concat, [1]] # cat head P4 45 | - [-1, 1, Conv, [256, 3, 1]] 46 | - [-1, 9, Conv, [256, 3, 1]] # 23 47 | 48 | - [-1, 1, Conv, [256, 3, 2]] 49 | - [[-1, 10], 1, Concat, [1]] # cat head P5 50 | - [-1, 1, Conv, [512, 3, 1]] 51 | - [-1, 9, Conv, [512, 3, 1]] # 27 52 | 53 | - [[19, 23, 27], 1, Detect, [nc]] # Detect(P3, P4, P5) 54 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/stnyolov8.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs 9 | s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs 10 | m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs 11 | l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs 12 | x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs 13 | 14 | # YOLOv8.0n backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [1024, True]] 26 | - [-1, 1, SPPF, [1024, 5]] # 9 27 | 28 | # YOLOv8.0n head 29 | head: 30 | - [-1, 1, SpatialTransformer, [3]] # Add STN layer here 31 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 32 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 33 | - [-1, 3, C2f, [512]] # 12 34 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 35 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 36 | - [-1, 3, C2f, [512]] # 12 37 | 38 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 39 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 40 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 41 | 42 | - [-1, 1, Conv, [256, 3, 2]] 43 | - [[-1, 12], 1, Concat, [1]] # cat head P4 44 | - [-1, 3, C2f, [512]] # 18 (P4/16-medium) 45 | 46 | - [-1, 1, Conv, [512, 3, 2]] 47 | - [[-1, 9], 1, Concat, [1]] # cat head P5 48 | - [-1, 3, C2f, [1024]] # 21 (P5/32-large) 49 | 50 | - [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5) 51 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/test.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs 9 | s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs 10 | m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs 11 | l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs 12 | x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs 13 | 14 | # YOLOv8.0n backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1,1,SpatialTransformer ,[3]] 18 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 19 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 20 | - [-1, 3, C2f, [128, True]] 21 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 22 | - [-1, 6, C2f, [256, True]] 23 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 24 | - [-1, 6, C2f, [512, True]] 25 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 26 | - [-1, 3, C2f, [1024, True]] 27 | - [-1, 1, SPPF, [1024, 5]] # 9 28 | 29 | # YOLOv8.0n head 30 | head: 31 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 32 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 33 | - [-1, 3, C2f, [512]] # 12 34 | 35 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 36 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 37 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 38 | 39 | - [-1, 1, Conv, [256, 3, 2]] 40 | - [[-1, 12], 1, Concat, [1]] # cat head P4 41 | - [-1, 3, C2f, [512]] # 18 (P4/16-medium) 42 | 43 | - [-1, 1, Conv, [512, 3, 2]] 44 | - [[-1, 9], 1, Concat, [1]] # cat head P5 45 | - [-1, 3, C2f, [1024]] # 21 (P5/32-large) 46 | 47 | - [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-cls.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8-cls image classification model. For Usage examples see https://docs.ultralytics.com/tasks/classify 3 | 4 | # Parameters 5 | nc: 1000 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 1024] 11 | l: [1.00, 1.00, 1024] 12 | x: [1.00, 1.25, 1024] 13 | 14 | # YOLOv8.0n backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [1024, True]] 26 | 27 | # YOLOv8.0n head 28 | head: 29 | - [-1, 1, Classify, [nc]] # Classify 30 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-p2.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8 object detection model with P2-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 768] 11 | l: [1.00, 1.00, 512] 12 | x: [1.00, 1.25, 512] 13 | 14 | # YOLOv8.0 backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [1024, True]] 26 | - [-1, 1, SPPF, [1024, 5]] # 9 27 | 28 | # YOLOv8.0-p2 head 29 | head: 30 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 31 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 32 | - [-1, 3, C2f, [512]] # 12 33 | 34 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 35 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 36 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 37 | 38 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 39 | - [[-1, 2], 1, Concat, [1]] # cat backbone P2 40 | - [-1, 3, C2f, [128]] # 18 (P2/4-xsmall) 41 | 42 | - [-1, 1, Conv, [128, 3, 2]] 43 | - [[-1, 15], 1, Concat, [1]] # cat head P3 44 | - [-1, 3, C2f, [256]] # 21 (P3/8-small) 45 | 46 | - [-1, 1, Conv, [256, 3, 2]] 47 | - [[-1, 12], 1, Concat, [1]] # cat head P4 48 | - [-1, 3, C2f, [512]] # 24 (P4/16-medium) 49 | 50 | - [-1, 1, Conv, [512, 3, 2]] 51 | - [[-1, 9], 1, Concat, [1]] # cat head P5 52 | - [-1, 3, C2f, [1024]] # 27 (P5/32-large) 53 | 54 | - [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P2, P3, P4, P5) 55 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-p6.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8 object detection model with P3-P6 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n-p6.yaml' will call yolov8-p6.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 768] 11 | l: [1.00, 1.00, 512] 12 | x: [1.00, 1.25, 512] 13 | 14 | # YOLOv8.0x6 backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [768, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [768, True]] 26 | - [-1, 1, Conv, [1024, 3, 2]] # 9-P6/64 27 | - [-1, 3, C2f, [1024, True]] 28 | - [-1, 1, SPPF, [1024, 5]] # 11 29 | 30 | # YOLOv8.0x6 head 31 | head: 32 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 33 | - [[-1, 8], 1, Concat, [1]] # cat backbone P5 34 | - [-1, 3, C2, [768, False]] # 14 35 | 36 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 37 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 38 | - [-1, 3, C2, [512, False]] # 17 39 | 40 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 41 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 42 | - [-1, 3, C2, [256, False]] # 20 (P3/8-small) 43 | 44 | - [-1, 1, Conv, [256, 3, 2]] 45 | - [[-1, 17], 1, Concat, [1]] # cat head P4 46 | - [-1, 3, C2, [512, False]] # 23 (P4/16-medium) 47 | 48 | - [-1, 1, Conv, [512, 3, 2]] 49 | - [[-1, 14], 1, Concat, [1]] # cat head P5 50 | - [-1, 3, C2, [768, False]] # 26 (P5/32-large) 51 | 52 | - [-1, 1, Conv, [768, 3, 2]] 53 | - [[-1, 11], 1, Concat, [1]] # cat head P6 54 | - [-1, 3, C2, [1024, False]] # 29 (P6/64-xlarge) 55 | 56 | - [[20, 23, 26, 29], 1, Detect, [nc]] # Detect(P3, P4, P5, P6) 57 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-pose-p6.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8-pose-p6 keypoints/pose estimation model. For Usage examples see https://docs.ultralytics.com/tasks/pose 3 | 4 | # Parameters 5 | nc: 1 # number of classes 6 | kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) 7 | scales: # model compound scaling constants, i.e. 'model=yolov8n-p6.yaml' will call yolov8-p6.yaml with scale 'n' 8 | # [depth, width, max_channels] 9 | n: [0.33, 0.25, 1024] 10 | s: [0.33, 0.50, 1024] 11 | m: [0.67, 0.75, 768] 12 | l: [1.00, 1.00, 512] 13 | x: [1.00, 1.25, 512] 14 | 15 | # YOLOv8.0x6 backbone 16 | backbone: 17 | # [from, repeats, module, args] 18 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 19 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 20 | - [-1, 3, C2f, [128, True]] 21 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 22 | - [-1, 6, C2f, [256, True]] 23 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 24 | - [-1, 6, C2f, [512, True]] 25 | - [-1, 1, Conv, [768, 3, 2]] # 7-P5/32 26 | - [-1, 3, C2f, [768, True]] 27 | - [-1, 1, Conv, [1024, 3, 2]] # 9-P6/64 28 | - [-1, 3, C2f, [1024, True]] 29 | - [-1, 1, SPPF, [1024, 5]] # 11 30 | 31 | # YOLOv8.0x6 head 32 | head: 33 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 34 | - [[-1, 8], 1, Concat, [1]] # cat backbone P5 35 | - [-1, 3, C2, [768, False]] # 14 36 | 37 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 38 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 39 | - [-1, 3, C2, [512, False]] # 17 40 | 41 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 42 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 43 | - [-1, 3, C2, [256, False]] # 20 (P3/8-small) 44 | 45 | - [-1, 1, Conv, [256, 3, 2]] 46 | - [[-1, 17], 1, Concat, [1]] # cat head P4 47 | - [-1, 3, C2, [512, False]] # 23 (P4/16-medium) 48 | 49 | - [-1, 1, Conv, [512, 3, 2]] 50 | - [[-1, 14], 1, Concat, [1]] # cat head P5 51 | - [-1, 3, C2, [768, False]] # 26 (P5/32-large) 52 | 53 | - [-1, 1, Conv, [768, 3, 2]] 54 | - [[-1, 11], 1, Concat, [1]] # cat head P6 55 | - [-1, 3, C2, [1024, False]] # 29 (P6/64-xlarge) 56 | 57 | - [[20, 23, 26, 29], 1, Pose, [nc, kpt_shape]] # Pose(P3, P4, P5, P6) 58 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-pose.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8-pose keypoints/pose estimation model. For Usage examples see https://docs.ultralytics.com/tasks/pose 3 | 4 | # Parameters 5 | nc: 1 # number of classes 6 | kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible) 7 | scales: # model compound scaling constants, i.e. 'model=yolov8n-pose.yaml' will call yolov8-pose.yaml with scale 'n' 8 | # [depth, width, max_channels] 9 | n: [0.33, 0.25, 1024] 10 | s: [0.33, 0.50, 1024] 11 | m: [0.67, 0.75, 768] 12 | l: [1.00, 1.00, 512] 13 | x: [1.00, 1.25, 512] 14 | 15 | # YOLOv8.0n backbone 16 | backbone: 17 | # [from, repeats, module, args] 18 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 19 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 20 | - [-1, 3, C2f, [128, True]] 21 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 22 | - [-1, 6, C2f, [256, True]] 23 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 24 | - [-1, 6, C2f, [512, True]] 25 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 26 | - [-1, 3, C2f, [1024, True]] 27 | - [-1, 1, SPPF, [1024, 5]] # 9 28 | 29 | # YOLOv8.0n head 30 | head: 31 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 32 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 33 | - [-1, 3, C2f, [512]] # 12 34 | 35 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 36 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 37 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 38 | 39 | - [-1, 1, Conv, [256, 3, 2]] 40 | - [[-1, 12], 1, Concat, [1]] # cat head P4 41 | - [-1, 3, C2f, [512]] # 18 (P4/16-medium) 42 | 43 | - [-1, 1, Conv, [512, 3, 2]] 44 | - [[-1, 9], 1, Concat, [1]] # cat head P5 45 | - [-1, 3, C2f, [1024]] # 21 (P5/32-large) 46 | 47 | - [[15, 18, 21], 1, Pose, [nc, kpt_shape]] # Pose(P3, P4, P5) 48 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-rtdetr.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs 9 | s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs 10 | m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs 11 | l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs 12 | x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs 13 | 14 | # YOLOv8.0n backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [1024, True]] 26 | - [-1, 1, SPPF, [1024, 5]] # 9 27 | 28 | # YOLOv8.0n head 29 | head: 30 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 31 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 32 | - [-1, 3, C2f, [512]] # 12 33 | 34 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 35 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 36 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 37 | 38 | - [-1, 1, Conv, [256, 3, 2]] 39 | - [[-1, 12], 1, Concat, [1]] # cat head P4 40 | - [-1, 3, C2f, [512]] # 18 (P4/16-medium) 41 | 42 | - [-1, 1, Conv, [512, 3, 2]] 43 | - [[-1, 9], 1, Concat, [1]] # cat head P5 44 | - [-1, 3, C2f, [1024]] # 21 (P5/32-large) 45 | 46 | - [[15, 18, 21], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5) 47 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-seg-p6.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8-seg-p6 instance segmentation model. For Usage examples see https://docs.ultralytics.com/tasks/segment 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n-seg-p6.yaml' will call yolov8-seg-p6.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 768] 11 | l: [1.00, 1.00, 512] 12 | x: [1.00, 1.25, 512] 13 | 14 | # YOLOv8.0x6 backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [768, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [768, True]] 26 | - [-1, 1, Conv, [1024, 3, 2]] # 9-P6/64 27 | - [-1, 3, C2f, [1024, True]] 28 | - [-1, 1, SPPF, [1024, 5]] # 11 29 | 30 | # YOLOv8.0x6 head 31 | head: 32 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 33 | - [[-1, 8], 1, Concat, [1]] # cat backbone P5 34 | - [-1, 3, C2, [768, False]] # 14 35 | 36 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 37 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 38 | - [-1, 3, C2, [512, False]] # 17 39 | 40 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 41 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 42 | - [-1, 3, C2, [256, False]] # 20 (P3/8-small) 43 | 44 | - [-1, 1, Conv, [256, 3, 2]] 45 | - [[-1, 17], 1, Concat, [1]] # cat head P4 46 | - [-1, 3, C2, [512, False]] # 23 (P4/16-medium) 47 | 48 | - [-1, 1, Conv, [512, 3, 2]] 49 | - [[-1, 14], 1, Concat, [1]] # cat head P5 50 | - [-1, 3, C2, [768, False]] # 26 (P5/32-large) 51 | 52 | - [-1, 1, Conv, [768, 3, 2]] 53 | - [[-1, 11], 1, Concat, [1]] # cat head P6 54 | - [-1, 3, C2, [1024, False]] # 29 (P6/64-xlarge) 55 | 56 | - [[20, 23, 26, 29], 1, Segment, [nc, 32, 256]] # Pose(P3, P4, P5, P6) 57 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8-seg.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8-seg instance segmentation model. For Usage examples see https://docs.ultralytics.com/tasks/segment 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n-seg.yaml' will call yolov8-seg.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] 9 | s: [0.33, 0.50, 1024] 10 | m: [0.67, 0.75, 768] 11 | l: [1.00, 1.00, 512] 12 | x: [1.00, 1.25, 512] 13 | 14 | # YOLOv8.0n backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 18 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 19 | - [-1, 3, C2f, [128, True]] 20 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 21 | - [-1, 6, C2f, [256, True]] 22 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 23 | - [-1, 6, C2f, [512, True]] 24 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 25 | - [-1, 3, C2f, [1024, True]] 26 | - [-1, 1, SPPF, [1024, 5]] # 9 27 | 28 | # YOLOv8.0n head 29 | head: 30 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 31 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 32 | - [-1, 3, C2f, [512]] # 12 33 | 34 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 35 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 36 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 37 | 38 | - [-1, 1, Conv, [256, 3, 2]] 39 | - [[-1, 12], 1, Concat, [1]] # cat head P4 40 | - [-1, 3, C2f, [512]] # 18 (P4/16-medium) 41 | 42 | - [-1, 1, Conv, [512, 3, 2]] 43 | - [[-1, 9], 1, Concat, [1]] # cat head P5 44 | - [-1, 3, C2f, [1024]] # 21 (P5/32-large) 45 | 46 | - [[15, 18, 21], 1, Segment, [nc, 32, 256]] # Segment(P3, P4, P5) 47 | -------------------------------------------------------------------------------- /ultralytics/cfg/models/v8/yolov8.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect 3 | 4 | # Parameters 5 | nc: 80 # number of classes 6 | scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' 7 | # [depth, width, max_channels] 8 | n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs 9 | s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs 10 | m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs 11 | l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs 12 | x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs 13 | 14 | # YOLOv8.0n backbone 15 | backbone: 16 | # [from, repeats, module, args] 17 | #- [-1,1,SpatialTransformer ,[3]] 18 | - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 19 | - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 20 | - [-1, 3, C2f, [128, True]] 21 | - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 22 | - [-1, 6, C2f, [256, True]] 23 | - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 24 | - [-1, 6, C2f, [512, True]] 25 | - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 26 | - [-1, 3, C2f, [1024, True]] 27 | - [-1, 1, SPPF, [1024, 5]] # 9 28 | 29 | # YOLOv8.0n head 30 | head: 31 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 32 | - [[-1, 6], 1, Concat, [1]] # cat backbone P4 33 | - [-1, 3, C2f, [512]] # 12 34 | 35 | - [-1, 1, nn.Upsample, [None, 2, 'nearest']] 36 | - [[-1, 4], 1, Concat, [1]] # cat backbone P3 37 | - [-1, 3, C2f, [256]] # 15 (P3/8-small) 38 | 39 | - [-1, 1, Conv, [256, 3, 2]] 40 | - [[-1, 12], 1, Concat, [1]] # cat head P4 41 | - [-1, 3, C2f, [512]] # 18 (P4/16-medium) 42 | 43 | - [-1, 1, Conv, [512, 3, 2]] 44 | - [[-1, 9], 1, Concat, [1]] # cat head P5 45 | - [-1, 3, C2f, [1024]] # 21 (P5/32-large) 46 | 47 | - [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5) -------------------------------------------------------------------------------- /ultralytics/cfg/trackers/botsort.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # Default YOLO tracker settings for BoT-SORT tracker https://github.com/NirAharon/BoT-SORT 3 | 4 | tracker_type: botsort # tracker type, ['botsort', 'bytetrack'] 5 | track_high_thresh: 0.5 # threshold for the first association 6 | track_low_thresh: 0.1 # threshold for the second association 7 | new_track_thresh: 0.6 # threshold for init new track if the detection does not match any tracks 8 | track_buffer: 30 # buffer to calculate the time when to remove tracks 9 | match_thresh: 0.8 # threshold for matching tracks 10 | # min_box_area: 10 # threshold for min box areas(for tracker evaluation, not used for now) 11 | # mot20: False # for tracker evaluation(not used for now) 12 | 13 | # BoT-SORT settings 14 | gmc_method: sparseOptFlow # method of global motion compensation 15 | # ReID model related thresh (not supported yet) 16 | proximity_thresh: 0.5 17 | appearance_thresh: 0.25 18 | with_reid: False 19 | -------------------------------------------------------------------------------- /ultralytics/cfg/trackers/bytetrack.yaml: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | # Default YOLO tracker settings for ByteTrack tracker https://github.com/ifzhang/ByteTrack 3 | 4 | tracker_type: bytetrack # tracker type, ['botsort', 'bytetrack'] 5 | track_high_thresh: 0.5 # threshold for the first association 6 | track_low_thresh: 0.1 # threshold for the second association 7 | new_track_thresh: 0.6 # threshold for init new track if the detection does not match any tracks 8 | track_buffer: 30 # buffer to calculate the time when to remove tracks 9 | match_thresh: 0.8 # threshold for matching tracks 10 | # min_box_area: 10 # threshold for min box areas(for tracker evaluation, not used for now) 11 | # mot20: False # for tracker evaluation(not used for now) 12 | -------------------------------------------------------------------------------- /ultralytics/data/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .base import BaseDataset 4 | from .build import build_dataloader, build_yolo_dataset, load_inference_source 5 | from .dataset import ClassificationDataset, SemanticDataset, YOLODataset 6 | 7 | __all__ = ('BaseDataset', 'ClassificationDataset', 'SemanticDataset', 'YOLODataset', 'build_yolo_dataset', 8 | 'build_dataloader', 'load_inference_source') 9 | -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/augment.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/augment.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/augment.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/augment.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/base.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/base.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/base.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/base.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/build.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/build.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/build.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/build.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/converter.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/converter.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/converter.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/converter.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/dataset.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/dataset.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/dataset.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/dataset.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/loaders.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/loaders.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/loaders.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/loaders.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/utils.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/utils.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/data/__pycache__/utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/data/__pycache__/utils.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/data/annotator.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from pathlib import Path 4 | 5 | from ultralytics import SAM, YOLO 6 | 7 | 8 | def auto_annotate(data, det_model='yolov8x.pt', sam_model='sam_b.pt', device='', output_dir=None): 9 | """ 10 | Automatically annotates images using a YOLO object detection model and a SAM segmentation model. 11 | 12 | Args: 13 | data (str): Path to a folder containing images to be annotated. 14 | det_model (str, optional): Pre-trained YOLO detection model. Defaults to 'yolov8x.pt'. 15 | sam_model (str, optional): Pre-trained SAM segmentation model. Defaults to 'sam_b.pt'. 16 | device (str, optional): Device to run the models on. Defaults to an empty string (CPU or GPU, if available). 17 | output_dir (str | None | optional): Directory to save the annotated results. 18 | Defaults to a 'labels' folder in the same directory as 'data'. 19 | 20 | Example: 21 | ```python 22 | from ultralytics.data.annotator import auto_annotate 23 | 24 | auto_annotate(data='ultralytics/assets', det_model='yolov8n.pt', sam_model='mobile_sam.pt') 25 | ``` 26 | """ 27 | det_model = YOLO(det_model) 28 | sam_model = SAM(sam_model) 29 | 30 | data = Path(data) 31 | if not output_dir: 32 | output_dir = data.parent / f'{data.stem}_auto_annotate_labels' 33 | Path(output_dir).mkdir(exist_ok=True, parents=True) 34 | 35 | det_results = det_model(data, stream=True, device=device) 36 | 37 | for result in det_results: 38 | class_ids = result.boxes.cls.int().tolist() # noqa 39 | if len(class_ids): 40 | boxes = result.boxes.xyxy # Boxes object for bbox outputs 41 | sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device) 42 | segments = sam_results[0].masks.xyn # noqa 43 | 44 | with open(f'{str(Path(output_dir) / Path(result.path).stem)}.txt', 'w') as f: 45 | for i in range(len(segments)): 46 | s = segments[i] 47 | if len(s) == 0: 48 | continue 49 | segment = map(str, segments[i].reshape(-1).tolist()) 50 | f.write(f'{class_ids[i]} ' + ' '.join(segment) + '\n') 51 | -------------------------------------------------------------------------------- /ultralytics/data/scripts/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Ultralytics YOLO 🚀, AGPL-3.0 license 3 | # Download latest models from https://github.com/ultralytics/assets/releases 4 | # Example usage: bash ultralytics/data/scripts/download_weights.sh 5 | # parent 6 | # └── weights 7 | # ├── yolov8n.pt ← downloads here 8 | # ├── yolov8s.pt 9 | # └── ... 10 | 11 | python - < w - threshold, 2] = w # x2 26 | boxes[boxes[:, 3] > h - threshold, 3] = h # y2 27 | return boxes 28 | 29 | 30 | def bbox_iou(box1, boxes, iou_thres=0.9, image_shape=(640, 640), raw_output=False): 31 | """ 32 | Compute the Intersection-Over-Union of a bounding box with respect to an array of other bounding boxes. 33 | 34 | Args: 35 | box1 (torch.Tensor): (4, ) 36 | boxes (torch.Tensor): (n, 4) 37 | iou_thres (float): IoU threshold 38 | image_shape (tuple): (height, width) 39 | raw_output (bool): If True, return the raw IoU values instead of the indices 40 | 41 | Returns: 42 | high_iou_indices (torch.Tensor): Indices of boxes with IoU > thres 43 | """ 44 | boxes = adjust_bboxes_to_image_border(boxes, image_shape) 45 | # Obtain coordinates for intersections 46 | x1 = torch.max(box1[0], boxes[:, 0]) 47 | y1 = torch.max(box1[1], boxes[:, 1]) 48 | x2 = torch.min(box1[2], boxes[:, 2]) 49 | y2 = torch.min(box1[3], boxes[:, 3]) 50 | 51 | # Compute the area of intersection 52 | intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0) 53 | 54 | # Compute the area of both individual boxes 55 | box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) 56 | box2_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) 57 | 58 | # Compute the area of union 59 | union = box1_area + box2_area - intersection 60 | 61 | # Compute the IoU 62 | iou = intersection / union # Should be shape (n, ) 63 | if raw_output: 64 | return 0 if iou.numel() == 0 else iou 65 | 66 | # return indices of boxes with IoU > thres 67 | return torch.nonzero(iou > iou_thres).flatten() 68 | -------------------------------------------------------------------------------- /ultralytics/models/fastsam/val.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.models.yolo.segment import SegmentationValidator 4 | from ultralytics.utils.metrics import SegmentMetrics 5 | 6 | 7 | class FastSAMValidator(SegmentationValidator): 8 | """ 9 | Custom validation class for fast SAM (Segment Anything Model) segmentation in Ultralytics YOLO framework. 10 | 11 | Extends the SegmentationValidator class, customizing the validation process specifically for fast SAM. This class 12 | sets the task to 'segment' and uses the SegmentMetrics for evaluation. Additionally, plotting features are disabled 13 | to avoid errors during validation. 14 | 15 | Attributes: 16 | dataloader: The data loader object used for validation. 17 | save_dir (str): The directory where validation results will be saved. 18 | pbar: A progress bar object. 19 | args: Additional arguments for customization. 20 | _callbacks: List of callback functions to be invoked during validation. 21 | """ 22 | 23 | def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None): 24 | """ 25 | Initialize the FastSAMValidator class, setting the task to 'segment' and metrics to SegmentMetrics. 26 | 27 | Args: 28 | dataloader (torch.utils.data.DataLoader): Dataloader to be used for validation. 29 | save_dir (Path, optional): Directory to save results. 30 | pbar (tqdm.tqdm): Progress bar for displaying progress. 31 | args (SimpleNamespace): Configuration for the validator. 32 | _callbacks (dict): Dictionary to store various callback functions. 33 | 34 | Notes: 35 | Plots for ConfusionMatrix and other related metrics are disabled in this class to avoid errors. 36 | """ 37 | super().__init__(dataloader, save_dir, pbar, args, _callbacks) 38 | self.args.task = 'segment' 39 | self.args.plots = False # disable ConfusionMatrix and other plots to avoid errors 40 | self.metrics = SegmentMetrics(save_dir=self.save_dir, on_plot=self.on_plot) 41 | -------------------------------------------------------------------------------- /ultralytics/models/nas/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .model import NAS 4 | from .predict import NASPredictor 5 | from .val import NASValidator 6 | 7 | __all__ = 'NASPredictor', 'NASValidator', 'NAS' 8 | -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/model.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/model.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/model.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/model.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/val.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/val.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/__pycache__/val.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/nas/__pycache__/val.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/nas/model.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | """ 3 | YOLO-NAS model interface. 4 | 5 | Example: 6 | ```python 7 | from ultralytics import NAS 8 | 9 | model = NAS('yolo_nas_s') 10 | results = model.predict('ultralytics/assets/bus.jpg') 11 | ``` 12 | """ 13 | 14 | from pathlib import Path 15 | 16 | import torch 17 | 18 | from ultralytics.engine.model import Model 19 | from ultralytics.utils.torch_utils import model_info, smart_inference_mode 20 | 21 | from .predict import NASPredictor 22 | from .val import NASValidator 23 | 24 | 25 | class NAS(Model): 26 | """ 27 | YOLO NAS model for object detection. 28 | 29 | This class provides an interface for the YOLO-NAS models and extends the `Model` class from Ultralytics engine. 30 | It is designed to facilitate the task of object detection using pre-trained or custom-trained YOLO-NAS models. 31 | 32 | Example: 33 | ```python 34 | from ultralytics import NAS 35 | 36 | model = NAS('yolo_nas_s') 37 | results = model.predict('ultralytics/assets/bus.jpg') 38 | ``` 39 | 40 | Attributes: 41 | model (str): Path to the pre-trained model or model name. Defaults to 'yolo_nas_s.pt'. 42 | 43 | Note: 44 | YOLO-NAS models only support pre-trained models. Do not provide YAML configuration files. 45 | """ 46 | 47 | def __init__(self, model='yolo_nas_s.pt') -> None: 48 | """Initializes the NAS model with the provided or default 'yolo_nas_s.pt' model.""" 49 | assert Path(model).suffix not in ('.yaml', '.yml'), 'YOLO-NAS models only support pre-trained models.' 50 | super().__init__(model, task='detect') 51 | 52 | @smart_inference_mode() 53 | def _load(self, weights: str, task: str): 54 | """Loads an existing NAS model weights or creates a new NAS model with pretrained weights if not provided.""" 55 | import super_gradients 56 | suffix = Path(weights).suffix 57 | if suffix == '.pt': 58 | self.model = torch.load(weights) 59 | elif suffix == '': 60 | self.model = super_gradients.training.models.get(weights, pretrained_weights='coco') 61 | # Standardize model 62 | self.model.fuse = lambda verbose=True: self.model 63 | self.model.stride = torch.tensor([32]) 64 | self.model.names = dict(enumerate(self.model._class_names)) 65 | self.model.is_fused = lambda: False # for info() 66 | self.model.yaml = {} # for info() 67 | self.model.pt_path = weights # for export() 68 | self.model.task = 'detect' # for export() 69 | 70 | def info(self, detailed=False, verbose=True): 71 | """ 72 | Logs model info. 73 | 74 | Args: 75 | detailed (bool): Show detailed information about model. 76 | verbose (bool): Controls verbosity. 77 | """ 78 | return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640) 79 | 80 | @property 81 | def task_map(self): 82 | """Returns a dictionary mapping tasks to respective predictor and validator classes.""" 83 | return {'detect': {'predictor': NASPredictor, 'validator': NASValidator}} 84 | -------------------------------------------------------------------------------- /ultralytics/models/nas/predict.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | import torch 4 | 5 | from ultralytics.engine.predictor import BasePredictor 6 | from ultralytics.engine.results import Results 7 | from ultralytics.utils import ops 8 | 9 | 10 | class NASPredictor(BasePredictor): 11 | """ 12 | Ultralytics YOLO NAS Predictor for object detection. 13 | 14 | This class extends the `BasePredictor` from Ultralytics engine and is responsible for post-processing the 15 | raw predictions generated by the YOLO NAS models. It applies operations like non-maximum suppression and 16 | scaling the bounding boxes to fit the original image dimensions. 17 | 18 | Attributes: 19 | args (Namespace): Namespace containing various configurations for post-processing. 20 | 21 | Example: 22 | ```python 23 | from ultralytics import NAS 24 | 25 | model = NAS('yolo_nas_s') 26 | predictor = model.predictor 27 | # Assumes that raw_preds, img, orig_imgs are available 28 | results = predictor.postprocess(raw_preds, img, orig_imgs) 29 | ``` 30 | 31 | Note: 32 | Typically, this class is not instantiated directly. It is used internally within the `NAS` class. 33 | """ 34 | 35 | def postprocess(self, preds_in, img, orig_imgs): 36 | """Postprocess predictions and returns a list of Results objects.""" 37 | 38 | # Cat boxes and class scores 39 | boxes = ops.xyxy2xywh(preds_in[0][0]) 40 | preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1) 41 | 42 | preds = ops.non_max_suppression(preds, 43 | self.args.conf, 44 | self.args.iou, 45 | agnostic=self.args.agnostic_nms, 46 | max_det=self.args.max_det, 47 | classes=self.args.classes) 48 | 49 | if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list 50 | orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) 51 | 52 | results = [] 53 | for i, pred in enumerate(preds): 54 | orig_img = orig_imgs[i] 55 | pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape) 56 | img_path = self.batch[0][i] 57 | results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred)) 58 | return results 59 | -------------------------------------------------------------------------------- /ultralytics/models/nas/val.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | import torch 4 | 5 | from ultralytics.models.yolo.detect import DetectionValidator 6 | from ultralytics.utils import ops 7 | 8 | __all__ = ['NASValidator'] 9 | 10 | 11 | class NASValidator(DetectionValidator): 12 | """ 13 | Ultralytics YOLO NAS Validator for object detection. 14 | 15 | Extends `DetectionValidator` from the Ultralytics models package and is designed to post-process the raw predictions 16 | generated by YOLO NAS models. It performs non-maximum suppression to remove overlapping and low-confidence boxes, 17 | ultimately producing the final detections. 18 | 19 | Attributes: 20 | args (Namespace): Namespace containing various configurations for post-processing, such as confidence and IoU thresholds. 21 | lb (torch.Tensor): Optional tensor for multilabel NMS. 22 | 23 | Example: 24 | ```python 25 | from ultralytics import NAS 26 | 27 | model = NAS('yolo_nas_s') 28 | validator = model.validator 29 | # Assumes that raw_preds are available 30 | final_preds = validator.postprocess(raw_preds) 31 | ``` 32 | 33 | Note: 34 | This class is generally not instantiated directly but is used internally within the `NAS` class. 35 | """ 36 | 37 | def postprocess(self, preds_in): 38 | """Apply Non-maximum suppression to prediction outputs.""" 39 | boxes = ops.xyxy2xywh(preds_in[0][0]) 40 | preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1) 41 | return ops.non_max_suppression(preds, 42 | self.args.conf, 43 | self.args.iou, 44 | labels=self.lb, 45 | multi_label=False, 46 | agnostic=self.args.single_cls, 47 | max_det=self.args.max_det, 48 | max_time_img=0.5) 49 | -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .model import RTDETR 4 | from .predict import RTDETRPredictor 5 | from .val import RTDETRValidator 6 | 7 | __all__ = 'RTDETRPredictor', 'RTDETRValidator', 'RTDETR' 8 | -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/model.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/model.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/model.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/model.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/train.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/train.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/train.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/train.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/val.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/val.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/__pycache__/val.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/rtdetr/__pycache__/val.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/model.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | """RT-DETR model interface.""" 3 | from ultralytics.engine.model import Model 4 | from ultralytics.nn.tasks import RTDETRDetectionModel 5 | 6 | from .predict import RTDETRPredictor 7 | from .train import RTDETRTrainer 8 | from .val import RTDETRValidator 9 | 10 | 11 | class RTDETR(Model): 12 | """RTDETR model interface.""" 13 | 14 | def __init__(self, model='rtdetr-l.pt') -> None: 15 | """Initializes the RTDETR model with the given model file, defaulting to 'rtdetr-l.pt'.""" 16 | if model and model.split('.')[-1] not in ('pt', 'yaml', 'yml'): 17 | raise NotImplementedError('RT-DETR only supports creating from *.pt file or *.yaml file.') 18 | super().__init__(model=model, task='detect') 19 | 20 | @property 21 | def task_map(self): 22 | """Returns a dictionary mapping task names to corresponding Ultralytics task classes for RTDETR model.""" 23 | return { 24 | 'detect': { 25 | 'predictor': RTDETRPredictor, 26 | 'validator': RTDETRValidator, 27 | 'trainer': RTDETRTrainer, 28 | 'model': RTDETRDetectionModel}} 29 | -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/predict.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | import torch 4 | 5 | from ultralytics.data.augment import LetterBox 6 | from ultralytics.engine.predictor import BasePredictor 7 | from ultralytics.engine.results import Results 8 | from ultralytics.utils import ops 9 | 10 | 11 | class RTDETRPredictor(BasePredictor): 12 | """ 13 | A class extending the BasePredictor class for prediction based on an RT-DETR detection model. 14 | 15 | Example: 16 | ```python 17 | from ultralytics.utils import ASSETS 18 | from ultralytics.models.rtdetr import RTDETRPredictor 19 | 20 | args = dict(model='rtdetr-l.pt', source=ASSETS) 21 | predictor = RTDETRPredictor(overrides=args) 22 | predictor.predict_cli() 23 | ``` 24 | """ 25 | 26 | def postprocess(self, preds, img, orig_imgs): 27 | """Postprocess predictions and returns a list of Results objects.""" 28 | nd = preds[0].shape[-1] 29 | bboxes, scores = preds[0].split((4, nd - 4), dim=-1) 30 | 31 | if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list 32 | orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) 33 | 34 | results = [] 35 | for i, bbox in enumerate(bboxes): # (300, 4) 36 | bbox = ops.xywh2xyxy(bbox) 37 | score, cls = scores[i].max(-1, keepdim=True) # (300, 1) 38 | idx = score.squeeze(-1) > self.args.conf # (300, ) 39 | if self.args.classes is not None: 40 | idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx 41 | pred = torch.cat([bbox, score, cls], dim=-1)[idx] # filter 42 | orig_img = orig_imgs[i] 43 | oh, ow = orig_img.shape[:2] 44 | pred[..., [0, 2]] *= ow 45 | pred[..., [1, 3]] *= oh 46 | img_path = self.batch[0][i] 47 | results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred)) 48 | return results 49 | 50 | def pre_transform(self, im): 51 | """ 52 | Pre-transform input image before inference. 53 | 54 | Args: 55 | im (List(np.ndarray)): (N, 3, h, w) for tensor, [(h, w, 3) x N] for list. 56 | 57 | Notes: The size must be square(640) and scaleFilled. 58 | 59 | Returns: 60 | (list): A list of transformed imgs. 61 | """ 62 | letterbox = LetterBox(self.imgsz, auto=False, scaleFill=True) 63 | return [letterbox(image=x) for x in im] 64 | -------------------------------------------------------------------------------- /ultralytics/models/rtdetr/train.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from copy import copy 4 | 5 | import torch 6 | 7 | from ultralytics.models.yolo.detect import DetectionTrainer 8 | from ultralytics.nn.tasks import RTDETRDetectionModel 9 | from ultralytics.utils import RANK, colorstr 10 | 11 | from .val import RTDETRDataset, RTDETRValidator 12 | 13 | 14 | class RTDETRTrainer(DetectionTrainer): 15 | """ 16 | A class extending the DetectionTrainer class for training based on an RT-DETR detection model. 17 | 18 | Notes: 19 | - F.grid_sample used in rt-detr does not support the `deterministic=True` argument. 20 | - AMP training can lead to NaN outputs and may produce errors during bipartite graph matching. 21 | 22 | Example: 23 | ```python 24 | from ultralytics.models.rtdetr.train import RTDETRTrainer 25 | 26 | args = dict(model='rtdetr-l.yaml', data='coco8.yaml', imgsz=640, epochs=3) 27 | trainer = RTDETRTrainer(overrides=args) 28 | trainer.train() 29 | ``` 30 | """ 31 | 32 | def get_model(self, cfg=None, weights=None, verbose=True): 33 | """Return a YOLO detection model.""" 34 | model = RTDETRDetectionModel(cfg, nc=self.data['nc'], verbose=verbose and RANK == -1) 35 | if weights: 36 | model.load(weights) 37 | return model 38 | 39 | def build_dataset(self, img_path, mode='val', batch=None): 40 | """ 41 | Build RTDETR Dataset. 42 | 43 | Args: 44 | img_path (str): Path to the folder containing images. 45 | mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode. 46 | batch (int, optional): Size of batches, this is for `rect`. Defaults to None. 47 | """ 48 | return RTDETRDataset( 49 | img_path=img_path, 50 | imgsz=self.args.imgsz, 51 | batch_size=batch, 52 | augment=mode == 'train', # no augmentation 53 | hyp=self.args, 54 | rect=False, # no rect 55 | cache=self.args.cache or None, 56 | prefix=colorstr(f'{mode}: '), 57 | data=self.data) 58 | 59 | def get_validator(self): 60 | """Returns a DetectionValidator for RTDETR model validation.""" 61 | self.loss_names = 'giou_loss', 'cls_loss', 'l1_loss' 62 | return RTDETRValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args)) 63 | 64 | def preprocess_batch(self, batch): 65 | """Preprocesses a batch of images by scaling and converting to float.""" 66 | batch = super().preprocess_batch(batch) 67 | bs = len(batch['img']) 68 | batch_idx = batch['batch_idx'] 69 | gt_bbox, gt_class = [], [] 70 | for i in range(bs): 71 | gt_bbox.append(batch['bboxes'][batch_idx == i].to(batch_idx.device)) 72 | gt_class.append(batch['cls'][batch_idx == i].to(device=batch_idx.device, dtype=torch.long)) 73 | return batch 74 | -------------------------------------------------------------------------------- /ultralytics/models/sam/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .model import SAM 4 | from .predict import Predictor 5 | 6 | # from .build import build_sam 7 | 8 | __all__ = 'SAM', 'Predictor' # tuple or list 9 | -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/amg.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/amg.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/amg.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/amg.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/build.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/build.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/build.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/build.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/model.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/model.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/model.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/model.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/model.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | """SAM model interface.""" 3 | 4 | from pathlib import Path 5 | 6 | from ultralytics.engine.model import Model 7 | from ultralytics.utils.torch_utils import model_info 8 | 9 | from .build import build_sam 10 | from .predict import Predictor 11 | 12 | 13 | class SAM(Model): 14 | """SAM model interface.""" 15 | 16 | def __init__(self, model='sam_b.pt') -> None: 17 | """Initializes the SAM model instance with the specified pre-trained model file.""" 18 | if model and Path(model).suffix not in ('.pt', '.pth'): 19 | raise NotImplementedError('SAM prediction requires pre-trained *.pt or *.pth model.') 20 | super().__init__(model=model, task='segment') 21 | 22 | def _load(self, weights: str, task=None): 23 | """Loads the provided weights into the SAM model.""" 24 | self.model = build_sam(weights) 25 | 26 | def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs): 27 | """Predicts and returns segmentation masks for given image or video source.""" 28 | overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024) 29 | kwargs.update(overrides) 30 | prompts = dict(bboxes=bboxes, points=points, labels=labels) 31 | return super().predict(source, stream, prompts=prompts, **kwargs) 32 | 33 | def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs): 34 | """Calls the 'predict' function with given arguments to perform object detection.""" 35 | return self.predict(source, stream, bboxes, points, labels, **kwargs) 36 | 37 | def info(self, detailed=False, verbose=True): 38 | """ 39 | Logs model info. 40 | 41 | Args: 42 | detailed (bool): Show detailed information about model. 43 | verbose (bool): Controls verbosity. 44 | """ 45 | return model_info(self.model, detailed=detailed, verbose=verbose) 46 | 47 | @property 48 | def task_map(self): 49 | """Returns a dictionary mapping the 'segment' task to its corresponding 'Predictor'.""" 50 | return {'segment': {'predictor': Predictor}} 51 | -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/decoders.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/decoders.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/decoders.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/decoders.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/encoders.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/encoders.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/encoders.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/encoders.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/sam.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/sam.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/sam.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/sam.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/tiny_encoder.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/tiny_encoder.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/tiny_encoder.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/tiny_encoder.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/transformer.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/transformer.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/__pycache__/transformer.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/sam/modules/__pycache__/transformer.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/sam/modules/sam.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | # Copyright (c) Meta Platforms, Inc. and affiliates. 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | 9 | from typing import List 10 | 11 | import torch 12 | from torch import nn 13 | 14 | from .decoders import MaskDecoder 15 | from .encoders import ImageEncoderViT, PromptEncoder 16 | 17 | 18 | class Sam(nn.Module): 19 | """ 20 | Sam (Segment Anything Model) is designed for object segmentation tasks. It uses image encoders to generate image 21 | embeddings, and prompt encoders to encode various types of input prompts. These embeddings are then used by the mask 22 | decoder to predict object masks. 23 | 24 | Attributes: 25 | mask_threshold (float): Threshold value for mask prediction. 26 | image_format (str): Format of the input image, default is 'RGB'. 27 | image_encoder (ImageEncoderViT): The backbone used to encode the image into embeddings. 28 | prompt_encoder (PromptEncoder): Encodes various types of input prompts. 29 | mask_decoder (MaskDecoder): Predicts object masks from the image and prompt embeddings. 30 | pixel_mean (List[float]): Mean pixel values for image normalization. 31 | pixel_std (List[float]): Standard deviation values for image normalization. 32 | """ 33 | mask_threshold: float = 0.0 34 | image_format: str = 'RGB' 35 | 36 | def __init__( 37 | self, 38 | image_encoder: ImageEncoderViT, 39 | prompt_encoder: PromptEncoder, 40 | mask_decoder: MaskDecoder, 41 | pixel_mean: List[float] = (123.675, 116.28, 103.53), 42 | pixel_std: List[float] = (58.395, 57.12, 57.375) 43 | ) -> None: 44 | """ 45 | Initialize the Sam class to predict object masks from an image and input prompts. 46 | 47 | Note: 48 | All forward() operations moved to SAMPredictor. 49 | 50 | Args: 51 | image_encoder (ImageEncoderViT): The backbone used to encode the image into image embeddings. 52 | prompt_encoder (PromptEncoder): Encodes various types of input prompts. 53 | mask_decoder (MaskDecoder): Predicts masks from the image embeddings and encoded prompts. 54 | pixel_mean (List[float], optional): Mean values for normalizing pixels in the input image. Defaults to 55 | (123.675, 116.28, 103.53). 56 | pixel_std (List[float], optional): Std values for normalizing pixels in the input image. Defaults to 57 | (58.395, 57.12, 57.375). 58 | """ 59 | super().__init__() 60 | self.image_encoder = image_encoder 61 | self.prompt_encoder = prompt_encoder 62 | self.mask_decoder = mask_decoder 63 | self.register_buffer('pixel_mean', torch.Tensor(pixel_mean).view(-1, 1, 1), False) 64 | self.register_buffer('pixel_std', torch.Tensor(pixel_std).view(-1, 1, 1), False) 65 | -------------------------------------------------------------------------------- /ultralytics/models/utils/__pycache__/loss.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/utils/__pycache__/loss.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/utils/__pycache__/ops.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/utils/__pycache__/ops.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.models.yolo import classify, detect, pose, segment 4 | 5 | from .model import YOLO 6 | 7 | __all__ = 'classify', 'segment', 'detect', 'pose', 'YOLO' 8 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/__pycache__/model.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/__pycache__/model.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/__pycache__/model.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/__pycache__/model.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.models.yolo.classify.predict import ClassificationPredictor 4 | from ultralytics.models.yolo.classify.train import ClassificationTrainer 5 | from ultralytics.models.yolo.classify.val import ClassificationValidator 6 | 7 | __all__ = 'ClassificationPredictor', 'ClassificationTrainer', 'ClassificationValidator' 8 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/train.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/train.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/train.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/train.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/val.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/val.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/__pycache__/val.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/classify/__pycache__/val.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/classify/predict.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | import torch 4 | 5 | from ultralytics.engine.predictor import BasePredictor 6 | from ultralytics.engine.results import Results 7 | from ultralytics.utils import DEFAULT_CFG, ops 8 | 9 | 10 | class ClassificationPredictor(BasePredictor): 11 | """ 12 | A class extending the BasePredictor class for prediction based on a classification model. 13 | 14 | Notes: 15 | - Torchvision classification models can also be passed to the 'model' argument, i.e. model='resnet18'. 16 | 17 | Example: 18 | ```python 19 | from ultralytics.utils import ASSETS 20 | from ultralytics.models.yolo.classify import ClassificationPredictor 21 | 22 | args = dict(model='yolov8n-cls.pt', source=ASSETS) 23 | predictor = ClassificationPredictor(overrides=args) 24 | predictor.predict_cli() 25 | ``` 26 | """ 27 | 28 | def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None): 29 | """Initializes ClassificationPredictor setting the task to 'classify'.""" 30 | super().__init__(cfg, overrides, _callbacks) 31 | self.args.task = 'classify' 32 | 33 | def preprocess(self, img): 34 | """Converts input image to model-compatible data type.""" 35 | if not isinstance(img, torch.Tensor): 36 | img = torch.stack([self.transforms(im) for im in img], dim=0) 37 | img = (img if isinstance(img, torch.Tensor) else torch.from_numpy(img)).to(self.model.device) 38 | return img.half() if self.model.fp16 else img.float() # uint8 to fp16/32 39 | 40 | def postprocess(self, preds, img, orig_imgs): 41 | """Post-processes predictions to return Results objects.""" 42 | if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list 43 | orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) 44 | 45 | results = [] 46 | for i, pred in enumerate(preds): 47 | orig_img = orig_imgs[i] 48 | img_path = self.batch[0][i] 49 | results.append(Results(orig_img, path=img_path, names=self.model.names, probs=pred)) 50 | return results 51 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .predict import DetectionPredictor 4 | from .train import DetectionTrainer 5 | from .val import DetectionValidator 6 | 7 | __all__ = 'DetectionPredictor', 'DetectionTrainer', 'DetectionValidator' 8 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/train.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/train.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/train.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/train.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/val.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/val.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/__pycache__/val.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/detect/__pycache__/val.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/detect/predict.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.engine.predictor import BasePredictor 4 | from ultralytics.engine.results import Results 5 | from ultralytics.utils import ops 6 | 7 | 8 | class DetectionPredictor(BasePredictor): 9 | """ 10 | A class extending the BasePredictor class for prediction based on a detection model. 11 | 12 | Example: 13 | ```python 14 | from ultralytics.utils import ASSETS 15 | from ultralytics.models.yolo.detect import DetectionPredictor 16 | 17 | args = dict(model='yolov8n.pt', source=ASSETS) 18 | predictor = DetectionPredictor(overrides=args) 19 | predictor.predict_cli() 20 | ``` 21 | """ 22 | 23 | def postprocess(self, preds, img, orig_imgs): 24 | """Post-processes predictions and returns a list of Results objects.""" 25 | preds = ops.non_max_suppression(preds, 26 | self.args.conf, 27 | self.args.iou, 28 | agnostic=self.args.agnostic_nms, 29 | max_det=self.args.max_det, 30 | classes=self.args.classes) 31 | 32 | if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list 33 | orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) 34 | 35 | results = [] 36 | for i, pred in enumerate(preds): 37 | orig_img = orig_imgs[i] 38 | pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape) 39 | img_path = self.batch[0][i] 40 | results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred)) 41 | return results 42 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/model.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.engine.model import Model 4 | from ultralytics.models import yolo # noqa 5 | from ultralytics.nn.tasks import ClassificationModel, DetectionModel, PoseModel, SegmentationModel 6 | 7 | 8 | class YOLO(Model): 9 | """YOLO (You Only Look Once) object detection model.""" 10 | 11 | @property 12 | def task_map(self): 13 | """Map head to model, trainer, validator, and predictor classes.""" 14 | return { 15 | 'classify': { 16 | 'model': ClassificationModel, 17 | 'trainer': yolo.classify.ClassificationTrainer, 18 | 'validator': yolo.classify.ClassificationValidator, 19 | 'predictor': yolo.classify.ClassificationPredictor, }, 20 | 'detect': { 21 | 'model': DetectionModel, 22 | 'trainer': yolo.detect.DetectionTrainer, 23 | 'validator': yolo.detect.DetectionValidator, 24 | 'predictor': yolo.detect.DetectionPredictor, }, 25 | 'segment': { 26 | 'model': SegmentationModel, 27 | 'trainer': yolo.segment.SegmentationTrainer, 28 | 'validator': yolo.segment.SegmentationValidator, 29 | 'predictor': yolo.segment.SegmentationPredictor, }, 30 | 'pose': { 31 | 'model': PoseModel, 32 | 'trainer': yolo.pose.PoseTrainer, 33 | 'validator': yolo.pose.PoseValidator, 34 | 'predictor': yolo.pose.PosePredictor, }, } 35 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .predict import PosePredictor 4 | from .train import PoseTrainer 5 | from .val import PoseValidator 6 | 7 | __all__ = 'PoseTrainer', 'PoseValidator', 'PosePredictor' 8 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/train.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/train.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/train.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/train.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/val.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/val.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/__pycache__/val.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/pose/__pycache__/val.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/predict.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.engine.results import Results 4 | from ultralytics.models.yolo.detect.predict import DetectionPredictor 5 | from ultralytics.utils import DEFAULT_CFG, LOGGER, ops 6 | 7 | 8 | class PosePredictor(DetectionPredictor): 9 | """ 10 | A class extending the DetectionPredictor class for prediction based on a pose model. 11 | 12 | Example: 13 | ```python 14 | from ultralytics.utils import ASSETS 15 | from ultralytics.models.yolo.pose import PosePredictor 16 | 17 | args = dict(model='yolov8n-pose.pt', source=ASSETS) 18 | predictor = PosePredictor(overrides=args) 19 | predictor.predict_cli() 20 | ``` 21 | """ 22 | 23 | def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None): 24 | """Initializes PosePredictor, sets task to 'pose' and logs a warning for using 'mps' as device.""" 25 | super().__init__(cfg, overrides, _callbacks) 26 | self.args.task = 'pose' 27 | if isinstance(self.args.device, str) and self.args.device.lower() == 'mps': 28 | LOGGER.warning("WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. " 29 | 'See https://github.com/ultralytics/ultralytics/issues/4031.') 30 | 31 | def postprocess(self, preds, img, orig_imgs): 32 | """Return detection results for a given input image or list of images.""" 33 | preds = ops.non_max_suppression(preds, 34 | self.args.conf, 35 | self.args.iou, 36 | agnostic=self.args.agnostic_nms, 37 | max_det=self.args.max_det, 38 | classes=self.args.classes, 39 | nc=len(self.model.names)) 40 | 41 | if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list 42 | orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) 43 | 44 | results = [] 45 | for i, pred in enumerate(preds): 46 | orig_img = orig_imgs[i] 47 | pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape).round() 48 | pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:] 49 | pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape) 50 | img_path = self.batch[0][i] 51 | results.append( 52 | Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts)) 53 | return results 54 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/pose/train.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from copy import copy 4 | 5 | from ultralytics.models import yolo 6 | from ultralytics.nn.tasks import PoseModel 7 | from ultralytics.utils import DEFAULT_CFG, LOGGER 8 | from ultralytics.utils.plotting import plot_images, plot_results 9 | 10 | 11 | class PoseTrainer(yolo.detect.DetectionTrainer): 12 | """ 13 | A class extending the DetectionTrainer class for training based on a pose model. 14 | 15 | Example: 16 | ```python 17 | from ultralytics.models.yolo.pose import PoseTrainer 18 | 19 | args = dict(model='yolov8n-pose.pt', data='coco8-pose.yaml', epochs=3) 20 | trainer = PoseTrainer(overrides=args) 21 | trainer.train() 22 | ``` 23 | """ 24 | 25 | def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None): 26 | """Initialize a PoseTrainer object with specified configurations and overrides.""" 27 | if overrides is None: 28 | overrides = {} 29 | overrides['task'] = 'pose' 30 | super().__init__(cfg, overrides, _callbacks) 31 | 32 | if isinstance(self.args.device, str) and self.args.device.lower() == 'mps': 33 | LOGGER.warning("WARNING ⚠️ Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. " 34 | 'See https://github.com/ultralytics/ultralytics/issues/4031.') 35 | 36 | def get_model(self, cfg=None, weights=None, verbose=True): 37 | """Get pose estimation model with specified configuration and weights.""" 38 | model = PoseModel(cfg, ch=3, nc=self.data['nc'], data_kpt_shape=self.data['kpt_shape'], verbose=verbose) 39 | if weights: 40 | model.load(weights) 41 | 42 | return model 43 | 44 | def set_model_attributes(self): 45 | """Sets keypoints shape attribute of PoseModel.""" 46 | super().set_model_attributes() 47 | self.model.kpt_shape = self.data['kpt_shape'] 48 | 49 | def get_validator(self): 50 | """Returns an instance of the PoseValidator class for validation.""" 51 | self.loss_names = 'box_loss', 'pose_loss', 'kobj_loss', 'cls_loss', 'dfl_loss' 52 | return yolo.pose.PoseValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args)) 53 | 54 | def plot_training_samples(self, batch, ni): 55 | """Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints.""" 56 | images = batch['img'] 57 | kpts = batch['keypoints'] 58 | cls = batch['cls'].squeeze(-1) 59 | bboxes = batch['bboxes'] 60 | paths = batch['im_file'] 61 | batch_idx = batch['batch_idx'] 62 | plot_images(images, 63 | batch_idx, 64 | cls, 65 | bboxes, 66 | kpts=kpts, 67 | paths=paths, 68 | fname=self.save_dir / f'train_batch{ni}.jpg', 69 | on_plot=self.on_plot) 70 | 71 | def plot_metrics(self): 72 | """Plots training/val metrics.""" 73 | plot_results(file=self.csv, pose=True, on_plot=self.on_plot) # save results.png 74 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .predict import SegmentationPredictor 4 | from .train import SegmentationTrainer 5 | from .val import SegmentationValidator 6 | 7 | __all__ = 'SegmentationPredictor', 'SegmentationTrainer', 'SegmentationValidator' 8 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/predict.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/predict.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/predict.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/predict.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/train.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/train.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/train.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/train.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/val.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/val.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/__pycache__/val.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/models/yolo/segment/__pycache__/val.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/predict.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.engine.results import Results 4 | from ultralytics.models.yolo.detect.predict import DetectionPredictor 5 | from ultralytics.utils import DEFAULT_CFG, ops 6 | 7 | 8 | class SegmentationPredictor(DetectionPredictor): 9 | """ 10 | A class extending the DetectionPredictor class for prediction based on a segmentation model. 11 | 12 | Example: 13 | ```python 14 | from ultralytics.utils import ASSETS 15 | from ultralytics.models.yolo.segment import SegmentationPredictor 16 | 17 | args = dict(model='yolov8n-seg.pt', source=ASSETS) 18 | predictor = SegmentationPredictor(overrides=args) 19 | predictor.predict_cli() 20 | ``` 21 | """ 22 | 23 | def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None): 24 | """Initializes the SegmentationPredictor with the provided configuration, overrides, and callbacks.""" 25 | super().__init__(cfg, overrides, _callbacks) 26 | self.args.task = 'segment' 27 | 28 | def postprocess(self, preds, img, orig_imgs): 29 | """Applies non-max suppression and processes detections for each image in an input batch.""" 30 | p = ops.non_max_suppression(preds[0], 31 | self.args.conf, 32 | self.args.iou, 33 | agnostic=self.args.agnostic_nms, 34 | max_det=self.args.max_det, 35 | nc=len(self.model.names), 36 | classes=self.args.classes) 37 | 38 | if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list 39 | orig_imgs = ops.convert_torch2numpy_batch(orig_imgs) 40 | 41 | results = [] 42 | proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported 43 | for i, pred in enumerate(p): 44 | orig_img = orig_imgs[i] 45 | img_path = self.batch[0][i] 46 | if not len(pred): # save empty boxes 47 | masks = None 48 | elif self.args.retina_masks: 49 | pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape) 50 | masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2]) # HWC 51 | else: 52 | masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True) # HWC 53 | pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape) 54 | results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks)) 55 | return results 56 | -------------------------------------------------------------------------------- /ultralytics/models/yolo/segment/train.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from copy import copy 4 | 5 | from ultralytics.models import yolo 6 | from ultralytics.nn.tasks import SegmentationModel 7 | from ultralytics.utils import DEFAULT_CFG, RANK 8 | from ultralytics.utils.plotting import plot_images, plot_results 9 | 10 | 11 | class SegmentationTrainer(yolo.detect.DetectionTrainer): 12 | """ 13 | A class extending the DetectionTrainer class for training based on a segmentation model. 14 | 15 | Example: 16 | ```python 17 | from ultralytics.models.yolo.segment import SegmentationTrainer 18 | 19 | args = dict(model='yolov8n-seg.pt', data='coco8-seg.yaml', epochs=3) 20 | trainer = SegmentationTrainer(overrides=args) 21 | trainer.train() 22 | ``` 23 | """ 24 | 25 | def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None): 26 | """Initialize a SegmentationTrainer object with given arguments.""" 27 | if overrides is None: 28 | overrides = {} 29 | overrides['task'] = 'segment' 30 | super().__init__(cfg, overrides, _callbacks) 31 | 32 | def get_model(self, cfg=None, weights=None, verbose=True): 33 | """Return SegmentationModel initialized with specified config and weights.""" 34 | model = SegmentationModel(cfg, ch=3, nc=self.data['nc'], verbose=verbose and RANK == -1) 35 | if weights: 36 | model.load(weights) 37 | 38 | return model 39 | 40 | def get_validator(self): 41 | """Return an instance of SegmentationValidator for validation of YOLO model.""" 42 | self.loss_names = 'box_loss', 'seg_loss', 'cls_loss', 'dfl_loss' 43 | return yolo.segment.SegmentationValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args)) 44 | 45 | def plot_training_samples(self, batch, ni): 46 | """Creates a plot of training sample images with labels and box coordinates.""" 47 | plot_images(batch['img'], 48 | batch['batch_idx'], 49 | batch['cls'].squeeze(-1), 50 | batch['bboxes'], 51 | batch['masks'], 52 | paths=batch['im_file'], 53 | fname=self.save_dir / f'train_batch{ni}.jpg', 54 | on_plot=self.on_plot) 55 | 56 | def plot_metrics(self): 57 | """Plots training/val metrics.""" 58 | plot_results(file=self.csv, segment=True, on_plot=self.on_plot) # save results.png 59 | -------------------------------------------------------------------------------- /ultralytics/nn/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .tasks import (BaseModel, ClassificationModel, DetectionModel, SegmentationModel, attempt_load_one_weight, 4 | attempt_load_weights, guess_model_scale, guess_model_task, parse_model, torch_safe_load, 5 | yaml_model_load) 6 | 7 | __all__ = ('attempt_load_one_weight', 'attempt_load_weights', 'parse_model', 'yaml_model_load', 'guess_model_task', 8 | 'guess_model_scale', 'torch_safe_load', 'DetectionModel', 'SegmentationModel', 'ClassificationModel', 9 | 'BaseModel') 10 | -------------------------------------------------------------------------------- /ultralytics/nn/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/__pycache__/autobackend.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/__pycache__/autobackend.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/__pycache__/autobackend.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/__pycache__/autobackend.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/__pycache__/tasks.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/__pycache__/tasks.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/__pycache__/tasks.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/__pycache__/tasks.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | """ 3 | Ultralytics modules. 4 | 5 | Example: 6 | Visualize a module with Netron. 7 | ```python 8 | from ultralytics.nn.modules import * 9 | import torch 10 | import os 11 | 12 | x = torch.ones(1, 128, 40, 40) 13 | m = Conv(128, 128) 14 | f = f'{m._get_name()}.onnx' 15 | torch.onnx.export(m, x, f) 16 | os.system(f'onnxsim {f} {f} && open {f}') 17 | ``` 18 | """ 19 | 20 | from .block import (SpatialTransformer, C1, C2, C3, C3TR, DFL, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x, GhostBottleneck, 21 | HGBlock, HGStem, Proto, RepC3) 22 | from .conv import (CBAM, ChannelAttention, Concat, Conv, Conv2, ConvTranspose, DWConv, DWConvTranspose2d, Focus, 23 | GhostConv, LightConv, RepConv, SpatialAttention) 24 | from .head import Classify, Detect, Pose, RTDETRDecoder, Segment 25 | from .transformer import (AIFI, MLP, DeformableTransformerDecoder, DeformableTransformerDecoderLayer, LayerNorm2d, 26 | MLPBlock, MSDeformAttn, TransformerBlock, TransformerEncoderLayer, TransformerLayer) 27 | 28 | __all__ = ('SpatialTransformer','Conv', 'Conv2', 'LightConv', 'RepConv', 'DWConv', 'DWConvTranspose2d', 'ConvTranspose', 'Focus', 29 | 'GhostConv', 'ChannelAttention', 'SpatialAttention', 'CBAM', 'Concat', 'TransformerLayer', 30 | 'TransformerBlock', 'MLPBlock', 'LayerNorm2d', 'DFL', 'HGBlock', 'HGStem', 'SPP', 'SPPF', 'C1', 'C2', 'C3', 31 | 'C2f', 'C3x', 'C3TR', 'C3Ghost', 'GhostBottleneck', 'Bottleneck', 'BottleneckCSP', 'Proto', 'Detect', 32 | 'Segment', 'Pose', 'Classify', 'TransformerEncoderLayer', 'RepC3', 'RTDETRDecoder', 'AIFI', 33 | 'DeformableTransformerDecoder', 'DeformableTransformerDecoderLayer', 'MSDeformAttn', 'MLP') 34 | -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/block.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/block.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/block.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/block.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/conv.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/conv.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/conv.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/conv.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/head.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/head.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/head.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/head.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/transformer.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/transformer.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/transformer.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/transformer.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/utils.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/utils.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/nn/modules/__pycache__/utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/nn/modules/__pycache__/utils.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/pipeline.png -------------------------------------------------------------------------------- /ultralytics/trackers/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .bot_sort import BOTSORT 4 | from .byte_tracker import BYTETracker 5 | from .track import register_tracker 6 | 7 | __all__ = 'register_tracker', 'BOTSORT', 'BYTETracker' # allow simpler import 8 | -------------------------------------------------------------------------------- /ultralytics/trackers/basetrack.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from collections import OrderedDict 4 | 5 | import numpy as np 6 | 7 | 8 | class TrackState: 9 | """Enumeration of possible object tracking states.""" 10 | 11 | New = 0 12 | Tracked = 1 13 | Lost = 2 14 | Removed = 3 15 | 16 | 17 | class BaseTrack: 18 | """Base class for object tracking, handling basic track attributes and operations.""" 19 | 20 | _count = 0 21 | 22 | track_id = 0 23 | is_activated = False 24 | state = TrackState.New 25 | 26 | history = OrderedDict() 27 | features = [] 28 | curr_feature = None 29 | score = 0 30 | start_frame = 0 31 | frame_id = 0 32 | time_since_update = 0 33 | 34 | # Multi-camera 35 | location = (np.inf, np.inf) 36 | 37 | @property 38 | def end_frame(self): 39 | """Return the last frame ID of the track.""" 40 | return self.frame_id 41 | 42 | @staticmethod 43 | def next_id(): 44 | """Increment and return the global track ID counter.""" 45 | BaseTrack._count += 1 46 | return BaseTrack._count 47 | 48 | def activate(self, *args): 49 | """Activate the track with the provided arguments.""" 50 | raise NotImplementedError 51 | 52 | def predict(self): 53 | """Predict the next state of the track.""" 54 | raise NotImplementedError 55 | 56 | def update(self, *args, **kwargs): 57 | """Update the track with new observations.""" 58 | raise NotImplementedError 59 | 60 | def mark_lost(self): 61 | """Mark the track as lost.""" 62 | self.state = TrackState.Lost 63 | 64 | def mark_removed(self): 65 | """Mark the track as removed.""" 66 | self.state = TrackState.Removed 67 | 68 | @staticmethod 69 | def reset_id(): 70 | """Reset the global track ID counter.""" 71 | BaseTrack._count = 0 72 | -------------------------------------------------------------------------------- /ultralytics/trackers/track.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from functools import partial 4 | 5 | import torch 6 | 7 | from ultralytics.utils import IterableSimpleNamespace, yaml_load 8 | from ultralytics.utils.checks import check_yaml 9 | 10 | from .bot_sort import BOTSORT 11 | from .byte_tracker import BYTETracker 12 | 13 | TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT} 14 | 15 | 16 | def on_predict_start(predictor, persist=False): 17 | """ 18 | Initialize trackers for object tracking during prediction. 19 | 20 | Args: 21 | predictor (object): The predictor object to initialize trackers for. 22 | persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False. 23 | 24 | Raises: 25 | AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'. 26 | """ 27 | if hasattr(predictor, 'trackers') and persist: 28 | return 29 | tracker = check_yaml(predictor.args.tracker) 30 | cfg = IterableSimpleNamespace(**yaml_load(tracker)) 31 | assert cfg.tracker_type in ['bytetrack', 'botsort'], \ 32 | f"Only support 'bytetrack' and 'botsort' for now, but got '{cfg.tracker_type}'" 33 | trackers = [] 34 | for _ in range(predictor.dataset.bs): 35 | tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30) 36 | trackers.append(tracker) 37 | predictor.trackers = trackers 38 | 39 | 40 | def on_predict_postprocess_end(predictor): 41 | """Postprocess detected boxes and update with object tracking.""" 42 | bs = predictor.dataset.bs 43 | im0s = predictor.batch[1] 44 | for i in range(bs): 45 | det = predictor.results[i].boxes.cpu().numpy() 46 | if len(det) == 0: 47 | continue 48 | tracks = predictor.trackers[i].update(det, im0s[i]) 49 | if len(tracks) == 0: 50 | continue 51 | idx = tracks[:, -1].astype(int) 52 | predictor.results[i] = predictor.results[i][idx] 53 | predictor.results[i].update(boxes=torch.as_tensor(tracks[:, :-1])) 54 | 55 | 56 | def register_tracker(model, persist): 57 | """ 58 | Register tracking callbacks to the model for object tracking during prediction. 59 | 60 | Args: 61 | model (object): The model object to register tracking callbacks for. 62 | persist (bool): Whether to persist the trackers if they already exist. 63 | """ 64 | model.add_callback('on_predict_start', partial(on_predict_start, persist=persist)) 65 | model.add_callback('on_predict_postprocess_end', on_predict_postprocess_end) 66 | -------------------------------------------------------------------------------- /ultralytics/trackers/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/autobatch.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/autobatch.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/autobatch.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/autobatch.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/checks.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/checks.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/checks.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/checks.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/dist.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/dist.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/dist.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/dist.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/downloads.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/downloads.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/downloads.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/downloads.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/files.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/files.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/files.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/files.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/instance.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/instance.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/instance.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/instance.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/loss.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/loss.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/loss.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/loss.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/metrics.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/metrics.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/metrics.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/metrics.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/ops.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/ops.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/ops.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/ops.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/patches.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/patches.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/patches.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/patches.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/plotting.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/plotting.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/plotting.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/plotting.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/tal.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/tal.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/tal.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/tal.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/torch_utils.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/torch_utils.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/__pycache__/torch_utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/__pycache__/torch_utils.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__init__.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from .base import add_integration_callbacks, default_callbacks, get_default_callbacks 4 | 5 | __all__ = 'add_integration_callbacks', 'default_callbacks', 'get_default_callbacks' 6 | -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/base.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/base.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/base.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/base.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/clearml.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/clearml.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/clearml.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/clearml.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/comet.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/comet.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/comet.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/comet.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/dvc.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/dvc.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/dvc.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/dvc.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/hub.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/hub.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/hub.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/hub.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/mlflow.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/mlflow.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/mlflow.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/mlflow.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/neptune.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/neptune.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/neptune.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/neptune.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/raytune.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/raytune.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/raytune.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/raytune.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/tensorboard.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/tensorboard.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/tensorboard.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/tensorboard.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/wb.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/wb.cpython-311.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/__pycache__/wb.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Advanced-Vision-and-Learning-Lab/STN-YOLO/0614a13145306ada7faf559d17f8a09796174306/ultralytics/utils/callbacks/__pycache__/wb.cpython-39.pyc -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/raytune.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.utils import SETTINGS 4 | 5 | try: 6 | assert SETTINGS['raytune'] is True # verify integration is enabled 7 | import ray 8 | from ray import tune 9 | from ray.air import session 10 | 11 | except (ImportError, AssertionError): 12 | tune = None 13 | 14 | 15 | def on_fit_epoch_end(trainer): 16 | """Sends training metrics to Ray Tune at end of each epoch.""" 17 | if ray.tune.is_session_enabled(): 18 | metrics = trainer.metrics 19 | metrics['epoch'] = trainer.epoch 20 | session.report(metrics) 21 | 22 | 23 | callbacks = { 24 | 'on_fit_epoch_end': on_fit_epoch_end, } if tune else {} 25 | -------------------------------------------------------------------------------- /ultralytics/utils/callbacks/tensorboard.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING, colorstr 4 | 5 | try: 6 | # WARNING: do not move import due to protobuf issue in https://github.com/ultralytics/ultralytics/pull/4674 7 | from torch.utils.tensorboard import SummaryWriter 8 | 9 | assert not TESTS_RUNNING # do not log pytest 10 | assert SETTINGS['tensorboard'] is True # verify integration is enabled 11 | WRITER = None # TensorBoard SummaryWriter instance 12 | 13 | except (ImportError, AssertionError, TypeError): 14 | # TypeError for handling 'Descriptors cannot not be created directly.' protobuf errors in Windows 15 | SummaryWriter = None 16 | 17 | 18 | def _log_scalars(scalars, step=0): 19 | """Logs scalar values to TensorBoard.""" 20 | if WRITER: 21 | for k, v in scalars.items(): 22 | WRITER.add_scalar(k, v, step) 23 | 24 | 25 | def _log_tensorboard_graph(trainer): 26 | """Log model graph to TensorBoard.""" 27 | try: 28 | import warnings 29 | 30 | from ultralytics.utils.torch_utils import de_parallel, torch 31 | 32 | imgsz = trainer.args.imgsz 33 | imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz 34 | p = next(trainer.model.parameters()) # for device, type 35 | im = torch.zeros((1, 3, *imgsz), device=p.device, dtype=p.dtype) # input image (must be zeros, not empty) 36 | with warnings.catch_warnings(): 37 | warnings.simplefilter('ignore', category=UserWarning) # suppress jit trace warning 38 | WRITER.add_graph(torch.jit.trace(de_parallel(trainer.model), im, strict=False), []) 39 | except Exception as e: 40 | LOGGER.warning(f'WARNING ⚠️ TensorBoard graph visualization failure {e}') 41 | 42 | 43 | def on_pretrain_routine_start(trainer): 44 | """Initialize TensorBoard logging with SummaryWriter.""" 45 | if SummaryWriter: 46 | try: 47 | global WRITER 48 | WRITER = SummaryWriter(str(trainer.save_dir)) 49 | prefix = colorstr('TensorBoard: ') 50 | LOGGER.info(f"{prefix}Start with 'tensorboard --logdir {trainer.save_dir}', view at http://localhost:6006/") 51 | except Exception as e: 52 | LOGGER.warning(f'WARNING ⚠️ TensorBoard not initialized correctly, not logging this run. {e}') 53 | 54 | 55 | def on_train_start(trainer): 56 | """Log TensorBoard graph.""" 57 | if WRITER: 58 | _log_tensorboard_graph(trainer) 59 | 60 | 61 | def on_batch_end(trainer): 62 | """Logs scalar statistics at the end of a training batch.""" 63 | _log_scalars(trainer.label_loss_items(trainer.tloss, prefix='train'), trainer.epoch + 1) 64 | 65 | 66 | def on_fit_epoch_end(trainer): 67 | """Logs epoch metrics at end of training epoch.""" 68 | _log_scalars(trainer.metrics, trainer.epoch + 1) 69 | 70 | 71 | callbacks = { 72 | 'on_pretrain_routine_start': on_pretrain_routine_start, 73 | 'on_train_start': on_train_start, 74 | 'on_fit_epoch_end': on_fit_epoch_end, 75 | 'on_batch_end': on_batch_end} if SummaryWriter else {} 76 | -------------------------------------------------------------------------------- /ultralytics/utils/dist.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | import os 4 | import re 5 | import shutil 6 | import socket 7 | import sys 8 | import tempfile 9 | from pathlib import Path 10 | 11 | from . import USER_CONFIG_DIR 12 | from .torch_utils import TORCH_1_9 13 | 14 | 15 | def find_free_network_port() -> int: 16 | """ 17 | Finds a free port on localhost. 18 | 19 | It is useful in single-node training when we don't want to connect to a real main node but have to set the 20 | `MASTER_PORT` environment variable. 21 | """ 22 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: 23 | s.bind(('127.0.0.1', 0)) 24 | return s.getsockname()[1] # port 25 | 26 | 27 | def generate_ddp_file(trainer): 28 | """Generates a DDP file and returns its file name.""" 29 | module, name = f'{trainer.__class__.__module__}.{trainer.__class__.__name__}'.rsplit('.', 1) 30 | 31 | content = f'''overrides = {vars(trainer.args)} \nif __name__ == "__main__": 32 | from {module} import {name} 33 | from ultralytics.utils import DEFAULT_CFG_DICT 34 | 35 | cfg = DEFAULT_CFG_DICT.copy() 36 | cfg.update(save_dir='') # handle the extra key 'save_dir' 37 | trainer = {name}(cfg=cfg, overrides=overrides) 38 | trainer.train()''' 39 | (USER_CONFIG_DIR / 'DDP').mkdir(exist_ok=True) 40 | with tempfile.NamedTemporaryFile(prefix='_temp_', 41 | suffix=f'{id(trainer)}.py', 42 | mode='w+', 43 | encoding='utf-8', 44 | dir=USER_CONFIG_DIR / 'DDP', 45 | delete=False) as file: 46 | file.write(content) 47 | return file.name 48 | 49 | 50 | def generate_ddp_command(world_size, trainer): 51 | """Generates and returns command for distributed training.""" 52 | import __main__ # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218 53 | if not trainer.resume: 54 | shutil.rmtree(trainer.save_dir) # remove the save_dir 55 | file = str(Path(sys.argv[0]).resolve()) 56 | safe_pattern = re.compile(r'^[a-zA-Z0-9_. /\\-]{1,128}$') # allowed characters and maximum of 100 characters 57 | if not (safe_pattern.match(file) and Path(file).exists() and file.endswith('.py')): # using CLI 58 | file = generate_ddp_file(trainer) 59 | dist_cmd = 'torch.distributed.run' if TORCH_1_9 else 'torch.distributed.launch' 60 | port = find_free_network_port() 61 | cmd = [sys.executable, '-m', dist_cmd, '--nproc_per_node', f'{world_size}', '--master_port', f'{port}', file] 62 | return cmd, file 63 | 64 | 65 | def ddp_cleanup(trainer, file): 66 | """Delete temp file if created.""" 67 | if f'{id(trainer)}.py' in file: # if temp_file suffix in file 68 | os.remove(file) 69 | -------------------------------------------------------------------------------- /ultralytics/utils/errors.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | 3 | from ultralytics.utils import emojis 4 | 5 | 6 | class HUBModelError(Exception): 7 | """ 8 | Custom exception class for handling errors related to model fetching in Ultralytics YOLO. 9 | 10 | This exception is raised when a requested model is not found or cannot be retrieved. 11 | The message is also processed to include emojis for better user experience. 12 | 13 | Attributes: 14 | message (str): The error message displayed when the exception is raised. 15 | 16 | Note: 17 | The message is automatically processed through the 'emojis' function from the 'ultralytics.utils' package. 18 | """ 19 | 20 | def __init__(self, message='Model not found. Please check model URL and try again.'): 21 | """Create an exception for when a model is not found.""" 22 | super().__init__(emojis(message)) 23 | -------------------------------------------------------------------------------- /ultralytics/utils/patches.py: -------------------------------------------------------------------------------- 1 | # Ultralytics YOLO 🚀, AGPL-3.0 license 2 | """Monkey patches to update/extend functionality of existing functions.""" 3 | 4 | from pathlib import Path 5 | 6 | import cv2 7 | import numpy as np 8 | import torch 9 | 10 | # OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------ 11 | _imshow = cv2.imshow # copy to avoid recursion errors 12 | 13 | 14 | def imread(filename: str, flags: int = cv2.IMREAD_COLOR): 15 | """ 16 | Read an image from a file. 17 | 18 | Args: 19 | filename (str): Path to the file to read. 20 | flags (int, optional): Flag that can take values of cv2.IMREAD_*. Defaults to cv2.IMREAD_COLOR. 21 | 22 | Returns: 23 | (np.ndarray): The read image. 24 | """ 25 | return cv2.imdecode(np.fromfile(filename, np.uint8), flags) 26 | 27 | 28 | def imwrite(filename: str, img: np.ndarray, params=None): 29 | """ 30 | Write an image to a file. 31 | 32 | Args: 33 | filename (str): Path to the file to write. 34 | img (np.ndarray): Image to write. 35 | params (list of ints, optional): Additional parameters. See OpenCV documentation. 36 | 37 | Returns: 38 | (bool): True if the file was written, False otherwise. 39 | """ 40 | try: 41 | cv2.imencode(Path(filename).suffix, img, params)[1].tofile(filename) 42 | return True 43 | except Exception: 44 | return False 45 | 46 | 47 | def imshow(winname: str, mat: np.ndarray): 48 | """ 49 | Displays an image in the specified window. 50 | 51 | Args: 52 | winname (str): Name of the window. 53 | mat (np.ndarray): Image to be shown. 54 | """ 55 | _imshow(winname.encode('unicode_escape').decode(), mat) 56 | 57 | 58 | # PyTorch functions ---------------------------------------------------------------------------------------------------- 59 | _torch_save = torch.save # copy to avoid recursion errors 60 | 61 | 62 | def torch_save(*args, **kwargs): 63 | """ 64 | Use dill (if exists) to serialize the lambda functions where pickle does not do this. 65 | 66 | Args: 67 | *args (tuple): Positional arguments to pass to torch.save. 68 | **kwargs (dict): Keyword arguments to pass to torch.save. 69 | """ 70 | try: 71 | import dill as pickle # noqa 72 | except ImportError: 73 | import pickle 74 | 75 | if 'pickle_module' not in kwargs: 76 | kwargs['pickle_module'] = pickle # noqa 77 | return _torch_save(*args, **kwargs) 78 | -------------------------------------------------------------------------------- /yoloV8N.slurm: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ##NECESSARY JOB SPECIFICATIONS 4 | #SBATCH --job-name=JobXAI #Set the job name to "JobExample4" 5 | #SBATCH --time=48:30:00 #Set the wall clock limit to 1hr and 30min 6 | #SBATCH --ntasks=8 #Request 1 task 7 | #SBATCH --mem=9000M #Request 2560MB (2.5GB) per node 8 | #SBATCH --output=yolov8nstn_modified.%j #Send stdout/err to "Example4Out.[jobID]" 9 | #SBATCH --gres=gpu:2 #Request 1 GPU per node can be 1 or 2 10 | #SBATCH --partition=gpu #Request the GPU partition/queue 11 | 12 | ##OPTIONAL JOB SPECIFICATIONS 13 | ##SBATCH --account=123456 #Set billing account to 123456 14 | ##SBATCH --mail-type=ALL #Send email on all job events 15 | ##SBATCH --mail-user=yashzambre@tamu.edu #Send all emails to email_address 16 | 17 | #First Executable Line 18 | 19 | cd /scratch/user/yashzambre/ultralytics; 20 | ml GCCcore/9.3.0 GCC/9.3.0 Python/3.8.2 CUDAcore/11.0.2 CUDA/11.0.2 cuDNN/8.0.5.39-CUDA-11.0.2 21 | module load Anaconda3/2022; 22 | source activate myenv1; 23 | export MKL_THREADING_LAYER=GNU; 24 | python demo_og.py --yaml '/scratch/user/yashzambre/ultralytics/augmentations/5/data.yaml' --model 'YOLOV8N' --num_epochs 100 --folder 'Saved_Models/augmentation_5_YOLO_100' 25 | --------------------------------------------------------------------------------