ultralytics 8.0.142__tar.gz → 8.0.144__tar.gz

{ultralytics-8.0.142/ultralytics.egg-info → ultralytics-8.0.144}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics
-Version: 8.0.142
+Version: 8.0.144
 Summary: Ultralytics YOLOv8 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Home-page: https://github.com/ultralytics/ultralytics
 Author: Ultralytics

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/__init__.py

@@ -1,10 +1,9 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = '8.0.142'
+__version__ = '8.0.144'
 
-from ultralytics.engine.model import YOLO
 from ultralytics.hub import start
-from ultralytics.models import RTDETR, SAM
+from ultralytics.models import RTDETR, SAM, YOLO
 from ultralytics.models.fastsam import FastSAM
 from ultralytics.models.nas import NAS
 from ultralytics.utils import SETTINGS as settings

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/engine/exporter.py

@@ -593,14 +593,43 @@ class Exporter:
         f_onnx, _ = self.export_onnx()
 
         # Export to TF
-        int8 = '-oiqt -qt per-tensor' if self.args.int8 else ''
-        cmd = f'onnx2tf -i "{f_onnx}" -o "{f}" -nuo --non_verbose {int8}'
-        LOGGER.info(f"\n{prefix} running '{cmd}'")
+        tmp_file = f / 'tmp_tflite_int8_calibration_images.npy'  # int8 calibration images file
+        if self.args.int8:
+            if self.args.data:
+                import numpy as np
+
+                from ultralytics.data.dataset import YOLODataset
+                from ultralytics.data.utils import check_det_dataset
+
+                # Generate calibration data for integer quantization
+                LOGGER.info(f"{prefix} collecting INT8 calibration images from 'data={self.args.data}'")
+                dataset = YOLODataset(check_det_dataset(self.args.data)['val'], imgsz=self.imgsz[0], augment=False)
+                images = []
+                n_images = 100  # maximum number of images
+                for n, batch in enumerate(dataset):
+                    if n >= n_images:
+                        break
+                    im = batch['img'].permute(1, 2, 0)[None]  # list to nparray, CHW to BHWC,
+                    images.append(im)
+                f.mkdir()
+                images = torch.cat(images, 0).float()
+                # mean = images.view(-1, 3).mean(0)  # imagenet mean [123.675, 116.28, 103.53]
+                # std = images.view(-1, 3).std(0)  # imagenet std [58.395, 57.12, 57.375]
+                np.save(str(tmp_file), images.numpy())  # BHWC
+                int8 = f'-oiqt -qt per-tensor -cind images "{tmp_file}" "[[[[0, 0, 0]]]]" "[[[[255, 255, 255]]]]"'
+            else:
+                int8 = '-oiqt -qt per-tensor'
+        else:
+            int8 = ''
+
+        cmd = f'onnx2tf -i "{f_onnx}" -o "{f}" -nuo --non_verbose {int8}'.strip()
+        LOGGER.info(f"{prefix} running '{cmd}'")
         subprocess.run(cmd, shell=True)
         yaml_save(f / 'metadata.yaml', self.metadata)  # add metadata.yaml
 
         # Remove/rename TFLite models
         if self.args.int8:
+            tmp_file.unlink(missing_ok=True)
             for file in f.rglob('*_dynamic_range_quant.tflite'):
                 file.rename(file.with_name(file.stem.replace('_dynamic_range_quant', '_int8') + file.suffix))
             for file in f.rglob('*_integer_quant_with_int16_act.tflite'):

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/engine/model.py

@@ -1,36 +1,24 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
+import inspect
 import sys
 from pathlib import Path
 from typing import Union
 
 from ultralytics.cfg import get_cfg
 from ultralytics.engine.exporter import Exporter
-from ultralytics.models import yolo  # noqa
-from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, PoseModel, SegmentationModel,
-                                  attempt_load_one_weight, guess_model_task, nn, yaml_model_load)
+from ultralytics.hub.utils import HUB_WEB_ROOT
+from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task, nn, yaml_model_load
 from ultralytics.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, ROOT, callbacks,
                                is_git_dir, yaml_load)
 from ultralytics.utils.checks import check_file, check_imgsz, check_pip_update_available, check_yaml
 from ultralytics.utils.downloads import GITHUB_ASSET_STEMS
 from ultralytics.utils.torch_utils import smart_inference_mode
 
-# Map head to model, trainer, validator, and predictor classes
-TASK_MAP = {
-    'classify': [
-        ClassificationModel, yolo.classify.ClassificationTrainer, yolo.classify.ClassificationValidator,
-        yolo.classify.ClassificationPredictor],
-    'detect':
-    [DetectionModel, yolo.detect.DetectionTrainer, yolo.detect.DetectionValidator, yolo.detect.DetectionPredictor],
-    'segment': [
-        SegmentationModel, yolo.segment.SegmentationTrainer, yolo.segment.SegmentationValidator,
-        yolo.segment.SegmentationPredictor],
-    'pose': [PoseModel, yolo.pose.PoseTrainer, yolo.pose.PoseValidator, yolo.pose.PosePredictor]}
 
-
-class YOLO:
+class Model:
     """
-    YOLO (You Only Look Once) object detection model.
+    A base model class to unify apis for all the models.
 
     Args:
         model (str, Path): Path to the model file to load or create.
@@ -81,13 +69,13 @@ class YOLO:
         self.predictor = None  # reuse predictor
         self.model = None  # model object
         self.trainer = None  # trainer object
-        self.task = None  # task type
         self.ckpt = None  # if loaded from *.pt
         self.cfg = None  # if loaded from *.yaml
         self.ckpt_path = None
         self.overrides = {}  # overrides for trainer object
         self.metrics = None  # validation/training metrics
         self.session = None  # HUB session
+        self.task = task  # task type
         model = str(model).strip()  # strip spaces
 
         # Check if Ultralytics HUB model from https://hub.ultralytics.com
@@ -109,32 +97,29 @@ class YOLO:
         """Calls the 'predict' function with given arguments to perform object detection."""
         return self.predict(source, stream, **kwargs)
 
-    def __getattr__(self, attr):
-        """Raises error if object has no requested attribute."""
-        name = self.__class__.__name__
-        raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
-
     @staticmethod
     def is_hub_model(model):
         """Check if the provided model is a HUB model."""
         return any((
-            model.startswith('https://hub.ultralytics.com/models/'),  # i.e. https://hub.ultralytics.com/models/MODEL_ID
+            model.startswith(f'{HUB_WEB_ROOT}/models/'),  # i.e. https://hub.ultralytics.com/models/MODEL_ID
             [len(x) for x in model.split('_')] == [42, 20],  # APIKEY_MODELID
             len(model) == 20 and not Path(model).exists() and all(x not in model for x in './\\')))  # MODELID
 
-    def _new(self, cfg: str, task=None, verbose=True):
+    def _new(self, cfg: str, task=None, model=None, verbose=True):
         """
         Initializes a new model and infers the task type from the model definitions.
 
         Args:
             cfg (str): model configuration file
             task (str | None): model task
+            model (BaseModel): Customized model.
             verbose (bool): display model info on load
         """
         cfg_dict = yaml_model_load(cfg)
         self.cfg = cfg
         self.task = task or guess_model_task(cfg_dict)
-        self.model = TASK_MAP[self.task][0](cfg_dict, verbose=verbose and RANK == -1)  # build model
+        model = model or self.smart_load('model')
+        self.model = model(cfg_dict, verbose=verbose and RANK == -1)  # build model
         self.overrides['model'] = self.cfg
 
         # Below added to allow export from yamls
@@ -217,7 +202,7 @@ class YOLO:
         self.model.fuse()
 
     @smart_inference_mode()
-    def predict(self, source=None, stream=False, **kwargs):
+    def predict(self, source=None, stream=False, predictor=None, **kwargs):
         """
         Perform prediction using the YOLO model.
 
@@ -225,6 +210,7 @@ class YOLO:
             source (str | int | PIL | np.ndarray): The source of the image to make predictions on.
                           Accepts all source types accepted by the YOLO model.
             stream (bool): Whether to stream the predictions or not. Defaults to False.
+            predictor (BasePredictor): Customized predictor.
             **kwargs : Additional keyword arguments passed to the predictor.
                        Check the 'configuration' section in the documentation for all available options.
 
@@ -236,6 +222,8 @@ class YOLO:
             LOGGER.warning(f"WARNING ⚠️ 'source' is missing. Using 'source={source}'.")
         is_cli = (sys.argv[0].endswith('yolo') or sys.argv[0].endswith('ultralytics')) and any(
             x in sys.argv for x in ('predict', 'track', 'mode=predict', 'mode=track'))
+        # Check prompts for SAM/FastSAM
+        prompts = kwargs.pop('prompts', None)
         overrides = self.overrides.copy()
         overrides['conf'] = 0.25
         overrides.update(kwargs)  # prefer kwargs
@@ -245,12 +233,16 @@ class YOLO:
             overrides['save'] = kwargs.get('save', False)  # do not save by default if called in Python
         if not self.predictor:
             self.task = overrides.get('task') or self.task
-            self.predictor = TASK_MAP[self.task][3](overrides=overrides, _callbacks=self.callbacks)
+            predictor = predictor or self.smart_load('predictor')
+            self.predictor = predictor(overrides=overrides, _callbacks=self.callbacks)
             self.predictor.setup_model(model=self.model, verbose=is_cli)
         else:  # only update args if predictor is already setup
             self.predictor.args = get_cfg(self.predictor.args, overrides)
             if 'project' in overrides or 'name' in overrides:
                 self.predictor.save_dir = self.predictor.get_save_dir()
+        # Set prompts for SAM/FastSAM
+        if len and hasattr(self.predictor, 'set_prompts'):
+            self.predictor.set_prompts(prompts)
         return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
 
     def track(self, source=None, stream=False, persist=False, **kwargs):
@@ -277,12 +269,13 @@ class YOLO:
         return self.predict(source=source, stream=stream, **kwargs)
 
     @smart_inference_mode()
-    def val(self, data=None, **kwargs):
+    def val(self, data=None, validator=None, **kwargs):
         """
         Validate a model on a given dataset.
 
         Args:
             data (str): The dataset to validate on. Accepts all formats accepted by yolo
+            validator (BaseValidator): Customized validator.
             **kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs
         """
         overrides = self.overrides.copy()
@@ -295,11 +288,12 @@ class YOLO:
             self.task = args.task
         else:
             args.task = self.task
+        validator = validator or self.smart_load('validator')
         if args.imgsz == DEFAULT_CFG.imgsz and not isinstance(self.model, (str, Path)):
             args.imgsz = self.model.args['imgsz']  # use trained imgsz unless custom value is passed
         args.imgsz = check_imgsz(args.imgsz, max_dim=1)
 
-        validator = TASK_MAP[self.task][2](args=args, _callbacks=self.callbacks)
+        validator = validator(args=args, _callbacks=self.callbacks)
         validator(model=self.model)
         self.metrics = validator.metrics
 
@@ -343,15 +337,18 @@ class YOLO:
             overrides['imgsz'] = self.model.args['imgsz']  # use trained imgsz unless custom value is passed
         if 'batch' not in kwargs:
             overrides['batch'] = 1  # default to 1 if not modified
+        if 'data' not in kwargs:
+            overrides['data'] = None  # default to None if not modified (avoid int8 calibration with coco.yaml)
         args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
         args.task = self.task
         return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model)
 
-    def train(self, **kwargs):
+    def train(self, trainer=None, **kwargs):
         """
         Trains the model on a given dataset.
 
         Args:
+            trainer (BaseTrainer, optional): Customized trainer.
             **kwargs (Any): Any number of arguments representing the training configuration.
         """
         self._check_is_pytorch_model()
@@ -371,7 +368,8 @@ class YOLO:
         if overrides.get('resume'):
             overrides['resume'] = self.ckpt_path
         self.task = overrides.get('task') or self.task
-        self.trainer = TASK_MAP[self.task][1](overrides=overrides, _callbacks=self.callbacks)
+        trainer = trainer or self.smart_load('trainer')
+        self.trainer = trainer(overrides=overrides, _callbacks=self.callbacks)
         if not overrides.get('resume'):  # manually set model only if not resuming
             self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
             self.model = self.trainer.model
@@ -440,3 +438,28 @@ class YOLO:
         """Reset all registered callbacks."""
         for event in callbacks.default_callbacks.keys():
             self.callbacks[event] = [callbacks.default_callbacks[event][0]]
+
+    def __getattr__(self, attr):
+        """Raises error if object has no requested attribute."""
+        name = self.__class__.__name__
+        raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
+
+    def smart_load(self, key):
+        """Load model/trainer/validator/predictor."""
+        try:
+            return self.task_map[self.task][key]
+        except Exception:
+            name = self.__class__.__name__
+            mode = inspect.stack()[1][3]  # get the function name.
+            raise NotImplementedError(
+                f'WARNING ⚠️ `{name}` model does not support `{mode}` mode for `{self.task}` task yet.')
+
+    @property
+    def task_map(self):
+        """
+        Map head to model, trainer, validator, and predictor classes.
+
+        Returns:
+            task_map (dict): The map of model task to mode classes.
+        """
+        raise NotImplementedError('Please provide task map for your model!')

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/engine/predictor.py

@@ -105,6 +105,7 @@ class BasePredictor:
         self.results = None
         self.transforms = None
         self.callbacks = _callbacks or callbacks.get_default_callbacks()
+        self.txt_path = None
         callbacks.add_integration_callbacks(self)
 
     def get_save_dir(self):
@@ -178,7 +179,8 @@ class BasePredictor:
         if self.args.save_txt:
             result.save_txt(f'{self.txt_path}.txt', save_conf=self.args.save_conf)
         if self.args.save_crop:
-            result.save_crop(save_dir=self.save_dir / 'crops', file_name=self.data_path.stem)
+            result.save_crop(save_dir=self.save_dir / 'crops',
+                             file_name=self.data_path.stem + ('' if self.dataset.mode == 'image' else f'_{frame}'))
 
         return log_string
 

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/engine/results.py

@@ -76,7 +76,6 @@ class Results(SimpleClass):
         probs (torch.tensor, optional): A 1D tensor of probabilities of each class for classification task.
         keypoints (List[List[float]], optional): A list of detected keypoints for each object.
 
-
     Attributes:
         orig_img (numpy.ndarray): The original image as a numpy array.
         orig_shape (tuple): The original image shape in (height, width) format.
@@ -172,6 +171,7 @@ class Results(SimpleClass):
             pil=False,
             img=None,
             im_gpu=None,
+            kpt_radius=5,
             kpt_line=True,
             labels=True,
             boxes=True,
@@ -190,6 +190,7 @@ class Results(SimpleClass):
             pil (bool): Whether to return the image as a PIL Image.
             img (numpy.ndarray): Plot to another image. if not, plot to original image.
             im_gpu (torch.Tensor): Normalized image in gpu with shape (1, 3, 640, 640), for faster mask plotting.
+            kpt_radius (int, optional): Radius of the drawn keypoints. Default is 5.
             kpt_line (bool): Whether to draw lines connecting keypoints.
             labels (bool): Whether to plot the label of bounding boxes.
             boxes (bool): Whether to plot the bounding boxes.
@@ -251,7 +252,7 @@ class Results(SimpleClass):
         # Plot Pose results
         if self.keypoints is not None:
             for k in reversed(self.keypoints.data):
-                annotator.kpts(k, self.orig_shape, kpt_line=kpt_line)
+                annotator.kpts(k, self.orig_shape, radius=kpt_radius, kpt_line=kpt_line)
 
         return annotator.result()
 

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/hub/__init__.py

@@ -4,7 +4,7 @@ import requests
 
 from ultralytics.data.utils import HUBDatasetStats
 from ultralytics.hub.auth import Auth
-from ultralytics.hub.utils import PREFIX
+from ultralytics.hub.utils import HUB_API_ROOT, HUB_WEB_ROOT, PREFIX
 from ultralytics.utils import LOGGER, SETTINGS, USER_CONFIG_DIR, yaml_save
 
 
@@ -50,13 +50,13 @@ WARNING ⚠️ ultralytics.start() is deprecated after 8.0.60. Updated usage to
 from ultralytics import YOLO, hub
 
 hub.login('{api_key}')
-model = YOLO('https://hub.ultralytics.com/models/{model_id}')
+model = YOLO('{HUB_WEB_ROOT}/models/{model_id}')
 model.train()""")
 
 
 def reset_model(model_id=''):
     """Reset a trained model to an untrained state."""
-    r = requests.post('https://api.ultralytics.com/model-reset', json={'apiKey': Auth().api_key, 'modelId': model_id})
+    r = requests.post(f'{HUB_API_ROOT}/model-reset', json={'apiKey': Auth().api_key, 'modelId': model_id})
     if r.status_code == 200:
         LOGGER.info(f'{PREFIX}Model reset successfully')
         return
@@ -72,7 +72,7 @@ def export_fmts_hub():
 def export_model(model_id='', format='torchscript'):
     """Export a model to all formats."""
     assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
-    r = requests.post(f'https://api.ultralytics.com/v1/models/{model_id}/export',
+    r = requests.post(f'{HUB_API_ROOT}/v1/models/{model_id}/export',
                       json={'format': format},
                       headers={'x-api-key': Auth().api_key})
     assert r.status_code == 200, f'{PREFIX}{format} export failure {r.status_code} {r.reason}'
@@ -82,7 +82,7 @@ def export_model(model_id='', format='torchscript'):
 def get_export(model_id='', format='torchscript'):
     """Get an exported model dictionary with download URL."""
     assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
-    r = requests.post('https://api.ultralytics.com/get-export',
+    r = requests.post(f'{HUB_API_ROOT}/get-export',
                       json={
                           'apiKey': Auth().api_key,
                           'modelId': model_id,
@@ -110,7 +110,7 @@ def check_dataset(path='', task='detect'):
         ```
     """
     HUBDatasetStats(path=path, task=task).get_json()
-    LOGGER.info('Checks completed correctly ✅. Upload this dataset to https://hub.ultralytics.com/datasets/.')
+    LOGGER.info(f'Checks completed correctly ✅. Upload this dataset to {HUB_WEB_ROOT}/datasets/.')
 
 
 if __name__ == '__main__':

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/hub/auth.py

@@ -2,10 +2,10 @@
 
 import requests
 
-from ultralytics.hub.utils import HUB_API_ROOT, PREFIX, request_with_credentials
+from ultralytics.hub.utils import HUB_API_ROOT, HUB_WEB_ROOT, PREFIX, request_with_credentials
 from ultralytics.utils import LOGGER, SETTINGS, emojis, is_colab
 
-API_KEY_URL = 'https://hub.ultralytics.com/settings?tab=api+keys'
+API_KEY_URL = f'{HUB_WEB_ROOT}/settings?tab=api+keys'
 
 
 class Auth:

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/hub/session.py

@@ -6,7 +6,7 @@ from time import sleep
 
 import requests
 
-from ultralytics.hub.utils import HUB_API_ROOT, PREFIX, smart_request
+from ultralytics.hub.utils import HUB_API_ROOT, HUB_WEB_ROOT, PREFIX, smart_request
 from ultralytics.utils import LOGGER, __version__, checks, emojis, is_colab, threaded
 from ultralytics.utils.errors import HUBModelError
 
@@ -49,21 +49,21 @@ class HUBTrainingSession:
         from ultralytics.hub.auth import Auth
 
         # Parse input
-        if url.startswith('https://hub.ultralytics.com/models/'):
-            url = url.split('https://hub.ultralytics.com/models/')[-1]
+        if url.startswith(f'{HUB_WEB_ROOT}/models/'):
+            url = url.split(f'{HUB_WEB_ROOT}/models/')[-1]
         if [len(x) for x in url.split('_')] == [42, 20]:
             key, model_id = url.split('_')
         elif len(url) == 20:
             key, model_id = '', url
         else:
             raise HUBModelError(f"model='{url}' not found. Check format is correct, i.e. "
-                                f"model='https://hub.ultralytics.com/models/MODEL_ID' and try again.")
+                                f"model='{HUB_WEB_ROOT}/models/MODEL_ID' and try again.")
 
         # Authorize
         auth = Auth(key)
         self.agent_id = None  # identifies which instance is communicating with server
         self.model_id = model_id
-        self.model_url = f'https://hub.ultralytics.com/models/{model_id}'
+        self.model_url = f'{HUB_WEB_ROOT}/models/{model_id}'
         self.api_url = f'{HUB_API_ROOT}/v1/models/{model_id}'
         self.auth_header = auth.get_auth_header()
         self.rate_limits = {'metrics': 3.0, 'ckpt': 900.0, 'heartbeat': 300.0}  # rate limits (seconds)

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/hub/utils.py

@@ -18,6 +18,7 @@ from ultralytics.utils.downloads import GITHUB_ASSET_NAMES
 PREFIX = colorstr('Ultralytics HUB: ')
 HELP_MSG = 'If this issue persists please visit https://github.com/ultralytics/hub/issues for assistance.'
 HUB_API_ROOT = os.environ.get('ULTRALYTICS_HUB_API', 'https://api.ultralytics.com')
+HUB_WEB_ROOT = os.environ.get('ULTRALYTICS_HUB_WEB', 'https://hub.ultralytics.com')
 
 
 def request_with_credentials(url: str) -> any:

ultralytics-8.0.144/ultralytics/models/__init__.py

@@ -0,0 +1,7 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from .rtdetr import RTDETR
+from .sam import SAM
+from .yolo import YOLO
+
+__all__ = 'YOLO', 'RTDETR', 'SAM'  # allow simpler import

ultralytics-8.0.144/ultralytics/models/fastsam/model.py

@@ -0,0 +1,31 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from pathlib import Path
+
+from ultralytics.engine.model import Model
+
+from .predict import FastSAMPredictor
+from .val import FastSAMValidator
+
+
+class FastSAM(Model):
+    """
+    FastSAM model interface.
+
+    Usage - Predict:
+        from ultralytics import FastSAM
+
+        model = FastSAM('last.pt')
+        results = model.predict('ultralytics/assets/bus.jpg')
+    """
+
+    def __init__(self, model='FastSAM-x.pt'):
+        """Call the __init__ method of the parent class (YOLO) with the updated default model"""
+        if model == 'FastSAM.pt':
+            model = 'FastSAM-x.pt'
+        assert Path(model).suffix != '.yaml', 'FastSAM models only support pre-trained models.'
+        super().__init__(model=model, task='segment')
+
+    @property
+    def task_map(self):
+        return {'segment': {'predictor': FastSAMPredictor, 'validator': FastSAMValidator}}

ultralytics-8.0.144/ultralytics/models/nas/model.py

@@ -0,0 +1,59 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+YOLO-NAS model interface.
+
+Usage - Predict:
+    from ultralytics import NAS
+
+    model = NAS('yolo_nas_s')
+    results = model.predict('ultralytics/assets/bus.jpg')
+"""
+
+from pathlib import Path
+
+import torch
+
+from ultralytics.engine.model import Model
+from ultralytics.utils.torch_utils import model_info, smart_inference_mode
+
+from .predict import NASPredictor
+from .val import NASValidator
+
+
+class NAS(Model):
+
+    def __init__(self, model='yolo_nas_s.pt') -> None:
+        assert Path(model).suffix != '.yaml', 'YOLO-NAS models only support pre-trained models.'
+        super().__init__(model, task='detect')
+
+    @smart_inference_mode()
+    def _load(self, weights: str, task: str):
+        # Load or create new NAS model
+        import super_gradients
+        suffix = Path(weights).suffix
+        if suffix == '.pt':
+            self.model = torch.load(weights)
+        elif suffix == '':
+            self.model = super_gradients.training.models.get(weights, pretrained_weights='coco')
+        # Standardize model
+        self.model.fuse = lambda verbose=True: self.model
+        self.model.stride = torch.tensor([32])
+        self.model.names = dict(enumerate(self.model._class_names))
+        self.model.is_fused = lambda: False  # for info()
+        self.model.yaml = {}  # for info()
+        self.model.pt_path = weights  # for export()
+        self.model.task = 'detect'  # for export()
+
+    def info(self, detailed=False, verbose=True):
+        """
+        Logs model info.
+
+        Args:
+            detailed (bool): Show detailed information about model.
+            verbose (bool): Controls verbosity.
+        """
+        return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)
+
+    @property
+    def task_map(self):
+        return {'detect': {'predictor': NASPredictor, 'validator': NASValidator}}

ultralytics-8.0.144/ultralytics/models/rtdetr/model.py

@@ -0,0 +1,30 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+RT-DETR model interface
+"""
+from ultralytics.engine.model import Model
+from ultralytics.nn.tasks import RTDETRDetectionModel
+
+from .predict import RTDETRPredictor
+from .train import RTDETRTrainer
+from .val import RTDETRValidator
+
+
+class RTDETR(Model):
+    """
+    RTDETR model interface.
+    """
+
+    def __init__(self, model='rtdetr-l.pt') -> None:
+        if model and not model.endswith('.pt') and not model.endswith('.yaml'):
+            raise NotImplementedError('RT-DETR only supports creating from pt file or yaml file.')
+        super().__init__(model=model, task='detect')
+
+    @property
+    def task_map(self):
+        return {
+            'detect': {
+                'predictor': RTDETRPredictor,
+                'validator': RTDETRValidator,
+                'trainer': RTDETRTrainer,
+                'model': RTDETRDetectionModel}}

{ultralytics-8.0.142 → ultralytics-8.0.144}/ultralytics/models/sam/model.py

@@ -3,51 +3,38 @@
 SAM model interface
 """
 
-from ultralytics.cfg import get_cfg
+from ultralytics.engine.model import Model
 from ultralytics.utils.torch_utils import model_info
 
 from .build import build_sam
 from .predict import Predictor
 
 
-class SAM:
+class SAM(Model):
+    """
+    SAM model interface.
+    """
 
     def __init__(self, model='sam_b.pt') -> None:
         if model and not model.endswith('.pt') and not model.endswith('.pth'):
             # Should raise AssertionError instead?
             raise NotImplementedError('Segment anything prediction requires pre-trained checkpoint')
-        self.model = build_sam(model)
-        self.task = 'segment'  # required
-        self.predictor = None  # reuse predictor
+        super().__init__(model=model, task='segment')
+
+    def _load(self, weights: str, task=None):
+        self.model = build_sam(weights)
 
     def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
         """Predicts and returns segmentation masks for given image or video source."""
         overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024)
-        overrides.update(kwargs)  # prefer kwargs
-        if not self.predictor:
-            self.predictor = Predictor(overrides=overrides)
-            self.predictor.setup_model(model=self.model)
-        else:  # only update args if predictor is already setup
-            self.predictor.args = get_cfg(self.predictor.args, overrides)
-        return self.predictor(source, stream=stream, bboxes=bboxes, points=points, labels=labels)
-
-    def train(self, **kwargs):
-        """Function trains models but raises an error as SAM models do not support training."""
-        raise NotImplementedError("SAM models don't support training")
-
-    def val(self, **kwargs):
-        """Run validation given dataset."""
-        raise NotImplementedError("SAM models don't support validation")
+        kwargs.update(overrides)
+        prompts = dict(bboxes=bboxes, points=points, labels=labels)
+        return super().predict(source, stream, prompts=prompts, **kwargs)
 
     def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
         """Calls the 'predict' function with given arguments to perform object detection."""
         return self.predict(source, stream, bboxes, points, labels, **kwargs)
 
-    def __getattr__(self, attr):
-        """Raises error if object has no requested attribute."""
-        name = self.__class__.__name__
-        raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
-
     def info(self, detailed=False, verbose=True):
         """
         Logs model info.
@@ -57,3 +44,7 @@ class SAM:
             verbose (bool): Controls verbosity.
         """
         return model_info(self.model, detailed=detailed, verbose=verbose)
+
+    @property
+    def task_map(self):
+        return {'segment': {'predictor': Predictor}}