Skip to content

Commit

Permalink
Accept model=yolov5*u.yaml arguments (ultralytics#4230)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored Aug 8, 2023
1 parent 7dcdca3 commit f5fc807
Show file tree
Hide file tree
Showing 10 changed files with 23 additions and 12 deletions.
4 changes: 1 addition & 3 deletions docs/models/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,10 @@ You can use many of these models directly in the Command Line Interface (CLI) or

## Usage

You can use RT-DETR for object detection tasks using the `ultralytics` pip package. The following is a sample code snippet showing how to use RT-DETR models for training and inference:
This example provides simple inference code for YOLO, SAM and RTDETR models. For more options including handling inference results see [Predict](../modes/predict.md) mode. For using models with additional modes see [Train](../modes/train.md), [Val](../modes/val.md) and [Export](../modes/export.md).

!!! example ""

This example provides simple inference code for YOLO, SAM and RTDETR models. For more options including handling inference results see [Predict](../modes/predict.md) mode. For using models with additional modes see [Train](../modes/train.md), [Val](../modes/val.md) and [Export](../modes/export.md).

=== "Python"

PyTorch pretrained `*.pt` models as well as configuration `*.yaml` files can be passed to the `YOLO()`, `SAM()`, `NAS()` and `RTDETR()` classes to create a model instance in python:
Expand Down
2 changes: 2 additions & 0 deletions ultralytics/data/annotator.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license

from pathlib import Path

from ultralytics import SAM, YOLO
Expand Down
2 changes: 2 additions & 0 deletions ultralytics/data/converter.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license

import json
from collections import defaultdict
from pathlib import Path
Expand Down
Empty file.
1 change: 1 addition & 0 deletions ultralytics/hub/session.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license

import signal
import sys
from pathlib import Path
Expand Down
1 change: 1 addition & 0 deletions ultralytics/utils/callbacks/dvc.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# Ultralytics YOLO 🚀, GPL-3.0 license

import os

import pkg_resources as pkg
Expand Down
1 change: 1 addition & 0 deletions ultralytics/utils/callbacks/wb.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license

from ultralytics.utils import SETTINGS, TESTS_RUNNING
from ultralytics.utils.torch_utils import model_info_for_loggers

Expand Down
22 changes: 13 additions & 9 deletions ultralytics/utils/checks.py
Original file line number Diff line number Diff line change
Expand Up @@ -341,15 +341,19 @@ def check_suffix(file='yolov8n.pt', suffix='.pt', msg=''):

def check_yolov5u_filename(file: str, verbose: bool = True):
"""Replace legacy YOLOv5 filenames with updated YOLOv5u filenames."""
if ('yolov3' in file or 'yolov5' in file) and 'u' not in file:
original_file = file
file = re.sub(r'(.*yolov5([nsmlx]))\.pt', '\\1u.pt', file) # i.e. yolov5n.pt -> yolov5nu.pt
file = re.sub(r'(.*yolov5([nsmlx])6)\.pt', '\\1u.pt', file) # i.e. yolov5n6.pt -> yolov5n6u.pt
file = re.sub(r'(.*yolov3(|-tiny|-spp))\.pt', '\\1u.pt', file) # i.e. yolov3-spp.pt -> yolov3-sppu.pt
if file != original_file and verbose:
LOGGER.info(f"PRO TIP 💡 Replace 'model={original_file}' with new 'model={file}'.\nYOLOv5 'u' models are "
f'trained with https://github.com/ultralytics/ultralytics and feature improved performance vs '
f'standard YOLOv5 models trained with https://github.com/ultralytics/yolov5.\n')
if 'yolov3' in file or 'yolov5' in file:
if 'u.yaml' in file:
file = file.replace('u.yaml', '.yaml') # i.e. yolov5nu.yaml -> yolov5n.yaml
elif '.pt' in file and 'u' not in file:
original_file = file
file = re.sub(r'(.*yolov5([nsmlx]))\.pt', '\\1u.pt', file) # i.e. yolov5n.pt -> yolov5nu.pt
file = re.sub(r'(.*yolov5([nsmlx])6)\.pt', '\\1u.pt', file) # i.e. yolov5n6.pt -> yolov5n6u.pt
file = re.sub(r'(.*yolov3(|-tiny|-spp))\.pt', '\\1u.pt', file) # i.e. yolov3-spp.pt -> yolov3-sppu.pt
if file != original_file and verbose:
LOGGER.info(
f"PRO TIP 💡 Replace 'model={original_file}' with new 'model={file}'.\nYOLOv5 'u' models are "
f'trained with https://github.com/ultralytics/ultralytics and feature improved performance vs '
f'standard YOLOv5 models trained with https://github.com/ultralytics/yolov5.\n')
return file


Expand Down
1 change: 1 addition & 0 deletions ultralytics/utils/torch_utils.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license

import math
import os
import platform
Expand Down
1 change: 1 addition & 0 deletions ultralytics/utils/tuner.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license

from ultralytics.cfg import TASK2DATA, TASK2METRIC
from ultralytics.utils import DEFAULT_CFG_DICT, LOGGER, NUM_THREADS

Expand Down

0 comments on commit f5fc807

Please sign in to comment.