init commit

ultralytics/models/nas/__init__.py

@@ -0,0 +1,7 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .model import NAS
+from .predict import NASPredictor
+from .val import NASValidator
+
+__all__ = "NASPredictor", "NASValidator", "NAS"

ultralytics/models/nas/model.py

@@ -0,0 +1,101 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from ultralytics.engine.model import Model
+from ultralytics.utils import DEFAULT_CFG_DICT
+from ultralytics.utils.downloads import attempt_download_asset
+from ultralytics.utils.patches import torch_load
+from ultralytics.utils.torch_utils import model_info
+
+from .predict import NASPredictor
+from .val import NASValidator
+
+
+class NAS(Model):
+    """
+    YOLO-NAS model for object detection.
+
+    This class provides an interface for the YOLO-NAS models and extends the `Model` class from ultralytics engine.
+    It is designed to facilitate the task of object detection using pre-trained or custom-trained YOLO-NAS models.
+
+    Attributes:
+        model (torch.nn.Module): The loaded YOLO-NAS model.
+        task (str): The task type for the model, defaults to 'detect'.
+        predictor (NASPredictor): The predictor instance for making predictions.
+        validator (NASValidator): The validator instance for model validation.
+
+    Methods:
+        info: Log model information and return model details.
+
+    Examples:
+        >>> from ultralytics import NAS
+        >>> model = NAS("yolo_nas_s")
+        >>> results = model.predict("ultralytics/assets/bus.jpg")
+
+    Notes:
+        YOLO-NAS models only support pre-trained models. Do not provide YAML configuration files.
+    """
+
+    def __init__(self, model: str = "yolo_nas_s.pt") -> None:
+        """Initialize the NAS model with the provided or default model."""
+        assert Path(model).suffix not in {".yaml", ".yml"}, "YOLO-NAS models only support pre-trained models."
+        super().__init__(model, task="detect")
+
+    def _load(self, weights: str, task=None) -> None:
+        """
+        Load an existing NAS model weights or create a new NAS model with pretrained weights.
+
+        Args:
+            weights (str): Path to the model weights file or model name.
+            task (str, optional): Task type for the model.
+        """
+        import super_gradients
+
+        suffix = Path(weights).suffix
+        if suffix == ".pt":
+            self.model = torch_load(attempt_download_asset(weights))
+        elif suffix == "":
+            self.model = super_gradients.training.models.get(weights, pretrained_weights="coco")
+
+        # Override the forward method to ignore additional arguments
+        def new_forward(x, *args, **kwargs):
+            """Ignore additional __call__ arguments."""
+            return self.model._original_forward(x)
+
+        self.model._original_forward = self.model.forward
+        self.model.forward = new_forward
+
+        # Standardize model attributes for compatibility
+        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()
+        self.model.args = {**DEFAULT_CFG_DICT, **self.overrides}  # for export()
+        self.model.eval()
+
+    def info(self, detailed: bool = False, verbose: bool = True) -> dict[str, Any]:
+        """
+        Log model information.
+
+        Args:
+            detailed (bool): Show detailed information about model.
+            verbose (bool): Controls verbosity.
+
+        Returns:
+            (dict[str, Any]): Model information dictionary.
+        """
+        return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)
+
+    @property
+    def task_map(self) -> dict[str, dict[str, Any]]:
+        """Return a dictionary mapping tasks to respective predictor and validator classes."""
+        return {"detect": {"predictor": NASPredictor, "validator": NASValidator}}

ultralytics/models/nas/predict.py

@@ -0,0 +1,58 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import torch
+
+from ultralytics.models.yolo.detect.predict import DetectionPredictor
+from ultralytics.utils import ops
+
+
+class NASPredictor(DetectionPredictor):
+    """
+    Ultralytics YOLO NAS Predictor for object detection.
+
+    This class extends the DetectionPredictor from ultralytics engine and is responsible for post-processing the
+    raw predictions generated by the YOLO NAS models. It applies operations like non-maximum suppression and
+    scaling the bounding boxes to fit the original image dimensions.
+
+    Attributes:
+        args (Namespace): Namespace containing various configurations for post-processing including confidence
+            threshold, IoU threshold, agnostic NMS flag, maximum detections, and class filtering options.
+        model (torch.nn.Module): The YOLO NAS model used for inference.
+        batch (list): Batch of inputs for processing.
+
+    Examples:
+        >>> from ultralytics import NAS
+        >>> model = NAS("yolo_nas_s")
+        >>> predictor = model.predictor
+
+        Assume that raw_preds, img, orig_imgs are available
+        >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
+
+    Notes:
+        Typically, this class is not instantiated directly. It is used internally within the NAS class.
+    """
+
+    def postprocess(self, preds_in, img, orig_imgs):
+        """
+        Postprocess NAS model predictions to generate final detection results.
+
+        This method takes raw predictions from a YOLO NAS model, converts bounding box formats, and applies
+        post-processing operations to generate the final detection results compatible with Ultralytics
+        result visualization and analysis tools.
+
+        Args:
+            preds_in (list): Raw predictions from the NAS model, typically containing bounding boxes and class scores.
+            img (torch.Tensor): Input image tensor that was fed to the model, with shape (B, C, H, W).
+            orig_imgs (list | torch.Tensor | np.ndarray): Original images before preprocessing, used for scaling
+                coordinates back to original dimensions.
+
+        Returns:
+            (list): List of Results objects containing the processed predictions for each image in the batch.
+
+        Examples:
+            >>> predictor = NAS("yolo_nas_s").predictor
+            >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
+        """
+        boxes = ops.xyxy2xywh(preds_in[0][0])  # Convert bounding boxes from xyxy to xywh format
+        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # Concatenate boxes with class scores
+        return super().postprocess(preds, img, orig_imgs)

ultralytics/models/nas/val.py

@@ -0,0 +1,39 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import torch
+
+from ultralytics.models.yolo.detect import DetectionValidator
+from ultralytics.utils import ops
+
+__all__ = ["NASValidator"]
+
+
+class NASValidator(DetectionValidator):
+    """
+    Ultralytics YOLO NAS Validator for object detection.
+
+    Extends DetectionValidator from the Ultralytics models package and is designed to post-process the raw predictions
+    generated by YOLO NAS models. It performs non-maximum suppression to remove overlapping and low-confidence boxes,
+    ultimately producing the final detections.
+
+    Attributes:
+        args (Namespace): Namespace containing various configurations for post-processing, such as confidence and IoU
+            thresholds.
+        lb (torch.Tensor): Optional tensor for multilabel NMS.
+
+    Examples:
+        >>> from ultralytics import NAS
+        >>> model = NAS("yolo_nas_s")
+        >>> validator = model.validator
+        >>> # Assumes that raw_preds are available
+        >>> final_preds = validator.postprocess(raw_preds)
+
+    Notes:
+        This class is generally not instantiated directly but is used internally within the NAS class.
+    """
+
+    def postprocess(self, preds_in):
+        """Apply Non-maximum suppression to prediction outputs."""
+        boxes = ops.xyxy2xywh(preds_in[0][0])  # Convert bounding box format from xyxy to xywh
+        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # Concatenate boxes with scores and permute
+        return super().postprocess(preds)