ultralytics 8.3.97__py3-none-any.whl → 8.3.99__py3-none-any.whl

ultralytics/models/yolo/model.py

@@ -4,7 +4,16 @@ from pathlib import Path
 
 from ultralytics.engine.model import Model
 from ultralytics.models import yolo
-from ultralytics.nn.tasks import ClassificationModel, DetectionModel, OBBModel, PoseModel, SegmentationModel, WorldModel
+from ultralytics.nn.tasks import (
+    ClassificationModel,
+    DetectionModel,
+    OBBModel,
+    PoseModel,
+    SegmentationModel,
+    WorldModel,
+    YOLOEModel,
+    YOLOESegModel,
+)
 from ultralytics.utils import ROOT, yaml_load
 
 
@@ -12,12 +21,16 @@ class YOLO(Model):
     """YOLO (You Only Look Once) object detection model."""
 
     def __init__(self, model="yolo11n.pt", task=None, verbose=False):
-        """Initialize YOLO model, switching to YOLOWorld if model filename contains '-world'."""
+        """Initialize YOLO model, switching to YOLOWorld/YOLOE if model filename contains '-world'/'yoloe'."""
         path = Path(model)
         if "-world" in path.stem and path.suffix in {".pt", ".yaml", ".yml"}:  # if YOLOWorld PyTorch model
             new_instance = YOLOWorld(path, verbose=verbose)
             self.__class__ = type(new_instance)
             self.__dict__ = new_instance.__dict__
+        elif "yoloe" in path.stem and path.suffix in {".pt", ".yaml", ".yml"}:  # if YOLOE PyTorch model
+            new_instance = YOLOE(path, task=task, verbose=verbose)
+            self.__class__ = type(new_instance)
+            self.__dict__ = new_instance.__dict__
         else:
             # Continue with default YOLO initialization
             super().__init__(model=model, task=task, verbose=verbose)
@@ -96,7 +109,7 @@ class YOLOWorld(Model):
         Set the model's class names for detection.
 
         Args:
-            classes (List(str)): A list of categories i.e. ["person"].
+            classes (list[str]): A list of categories i.e. ["person"].
         """
         self.model.set_classes(classes)
         # Remove background if it's given
@@ -108,3 +121,169 @@ class YOLOWorld(Model):
         # Reset method class names
         if self.predictor:
             self.predictor.model.names = classes
+
+
+class YOLOE(Model):
+    """YOLOE object detection and segmentation model."""
+
+    def __init__(self, model="yoloe-v8s-seg.pt", task=None, verbose=False) -> None:
+        """
+        Initialize YOLOE model with a pre-trained model file.
+
+        Args:
+            model (str | Path): Path to the pre-trained model file. Supports *.pt and *.yaml formats.
+            task (str, optional): Task type for the model. Auto-detected if None.
+            verbose (bool): If True, prints additional information during initialization.
+        """
+        super().__init__(model=model, task=task, verbose=verbose)
+
+        # Assign default COCO class names when there are no custom names
+        if not hasattr(self.model, "names"):
+            self.model.names = yaml_load(ROOT / "cfg/datasets/coco8.yaml").get("names")
+
+    @property
+    def task_map(self):
+        """Map head to model, validator, and predictor classes."""
+        return {
+            "detect": {
+                "model": YOLOEModel,
+                "validator": yolo.yoloe.YOLOEDetectValidator,
+                "predictor": yolo.detect.DetectionPredictor,
+                "trainer": yolo.yoloe.YOLOETrainer,
+            },
+            "segment": {
+                "model": YOLOESegModel,
+                "validator": yolo.yoloe.YOLOESegValidator,
+                "predictor": yolo.segment.SegmentationPredictor,
+                "trainer": yolo.yoloe.YOLOESegTrainer,
+            },
+        }
+
+    def get_text_pe(self, texts):
+        """Get text positional embeddings for the given texts."""
+        assert isinstance(self.model, YOLOEModel)
+        return self.model.get_text_pe(texts)
+
+    def get_visual_pe(self, img, visual):
+        """Get visual positional embeddings for the given image and visual features."""
+        assert isinstance(self.model, YOLOEModel)
+        return self.model.get_visual_pe(img, visual)
+
+    def set_vocab(self, vocab, names):
+        """Set vocabulary and class names for the model."""
+        assert isinstance(self.model, YOLOEModel)
+        self.model.set_vocab(vocab, names=names)
+
+    def get_vocab(self, names):
+        """Get vocabulary for the given class names."""
+        assert isinstance(self.model, YOLOEModel)
+        return self.model.get_vocab(names)
+
+    def set_classes(self, classes, embeddings):
+        """
+        Set the model's class names and embeddings for detection.
+
+        Args:
+            classes (list[str]): A list of categories i.e. ["person"].
+            embeddings (torch.Tensor): Embeddings corresponding to the classes.
+        """
+        assert isinstance(self.model, YOLOEModel)
+        self.model.set_classes(classes, embeddings)
+        # Verify no background class is present
+        assert " " not in classes
+        self.model.names = classes
+
+        # Reset method class names
+        if self.predictor:
+            self.predictor.model.names = classes
+
+    def val(
+        self,
+        validator=None,
+        load_vp=False,
+        refer_data=None,
+        **kwargs,
+    ):
+        """
+        Validate the model using text or visual prompts.
+
+        Args:
+            validator (callable, optional): A callable validator function. If None, a default validator is loaded.
+            load_vp (bool): Whether to load visual prompts. If False, text prompts are used.
+            refer_data (str, optional): Path to the reference data for visual prompts.
+            **kwargs (Any): Additional keyword arguments to override default settings.
+
+        Returns:
+            (dict): Validation statistics containing metrics computed during validation.
+        """
+        custom = {"rect": not load_vp}  # method defaults
+        args = {**self.overrides, **custom, **kwargs, "mode": "val"}  # highest priority args on the right
+
+        validator = (validator or self._smart_load("validator"))(args=args, _callbacks=self.callbacks)
+        validator(model=self.model, load_vp=load_vp, refer_data=refer_data)
+        self.metrics = validator.metrics
+        return validator.metrics
+
+    def predict(
+        self,
+        source=None,
+        stream: bool = False,
+        visual_prompts: dict = {},
+        refer_image=None,
+        predictor=None,
+        **kwargs,
+    ):
+        """
+        Run prediction on images, videos, directories, streams, etc.
+
+        Args:
+            source (str | int | PIL.Image | np.ndarray, optional): Source for prediction. Accepts image paths,
+                directory paths, URL/YouTube streams, PIL images, numpy arrays, or webcam indices.
+            stream (bool): Whether to stream the prediction results. If True, results are yielded as a
+                generator as they are computed.
+            visual_prompts (dict): Dictionary containing visual prompts for the model. Must include 'bboxes' and
+                'cls' keys when non-empty.
+            refer_image (str | PIL.Image | np.ndarray, optional): Reference image for visual prompts.
+            predictor (callable, optional): Custom predictor function. If None, a predictor is automatically
+                loaded based on the task.
+            **kwargs (Any): Additional keyword arguments passed to the predictor.
+
+        Returns:
+            (List | generator): List of Results objects or generator of Results objects if stream=True.
+
+        Examples:
+            >>> model = YOLOE("yoloe-v8s-seg.pt")
+            >>> results = model.predict("path/to/image.jpg")
+            >>> # With visual prompts
+            >>> prompts = {"bboxes": [[10, 20, 100, 200]], "cls": ["person"]}
+            >>> results = model.predict("path/to/image.jpg", visual_prompts=prompts)
+        """
+        if len(visual_prompts):
+            assert "bboxes" in visual_prompts and "cls" in visual_prompts, (
+                f"Expected 'bboxes' and 'cls' in visual prompts, but got {visual_prompts.keys()}"
+            )
+            assert len(visual_prompts["bboxes"]) == len(visual_prompts["cls"]), (
+                f"Expected equal number of bounding boxes and classes, but got {len(visual_prompts['bboxes'])} and "
+                f"{len(visual_prompts['cls'])} respectively"
+            )
+        self.predictor = (predictor or self._smart_load("predictor"))(
+            overrides={"task": "segment", "mode": "predict", "save": False, "verbose": False}, _callbacks=self.callbacks
+        )
+
+        if len(visual_prompts):
+            num_cls = (
+                max(len(set(c)) for c in visual_prompts["cls"])
+                if isinstance(source, list)  # means multiple images
+                else len(set(visual_prompts["cls"]))
+            )
+            self.model.model[-1].nc = num_cls
+            self.model.names = [f"object{i}" for i in range(num_cls)]
+            self.predictor.set_prompts(visual_prompts)
+
+        self.predictor.setup_model(model=self.model)
+        if refer_image is not None and len(visual_prompts):
+            vpe = self.predictor.get_vpe(refer_image)
+            self.model.set_classes(self.model.names, vpe)
+            self.predictor = None  # reset predictor
+
+        return super().predict(source, stream, **kwargs)

ultralytics/models/yolo/segment/val.py

@@ -364,29 +364,56 @@ class SegmentationValidator(DetectionValidator):
 
     def eval_json(self, stats):
         """Return COCO-style object detection evaluation metrics."""
-        if self.args.save_json and self.is_coco and len(self.jdict):
-            anno_json = self.data["path"] / "annotations/instances_val2017.json"  # annotations
+        if self.args.save_json and (self.is_lvis or self.is_coco) and len(self.jdict):
             pred_json = self.save_dir / "predictions.json"  # predictions
-            LOGGER.info(f"\nEvaluating pycocotools mAP using {pred_json} and {anno_json}...")
-            try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
-                check_requirements("pycocotools>=2.0.6")
-                from pycocotools.coco import COCO  # noqa
-                from pycocotools.cocoeval import COCOeval  # noqa
 
+            anno_json = (
+                self.data["path"]
+                / "annotations"
+                / ("instances_val2017.json" if self.is_coco else f"lvis_v1_{self.args.split}.json")
+            )  # annotations
+
+            pkg = "pycocotools" if self.is_coco else "lvis"
+            LOGGER.info(f"\nEvaluating {pkg} mAP using {pred_json} and {anno_json}...")
+            try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
                 for x in anno_json, pred_json:
                     assert x.is_file(), f"{x} file not found"
-                anno = COCO(str(anno_json))  # init annotations api
-                pred = anno.loadRes(str(pred_json))  # init predictions api (must pass string, not Path)
-                for i, eval in enumerate([COCOeval(anno, pred, "bbox"), COCOeval(anno, pred, "segm")]):
-                    if self.is_coco:
-                        eval.params.imgIds = [int(Path(x).stem) for x in self.dataloader.dataset.im_files]  # im to eval
+                check_requirements("pycocotools>=2.0.6" if self.is_coco else "lvis>=0.5.3")
+                if self.is_coco:
+                    from pycocotools.coco import COCO  # noqa
+                    from pycocotools.cocoeval import COCOeval  # noqa
+
+                    anno = COCO(str(anno_json))  # init annotations api
+                    pred = anno.loadRes(str(pred_json))  # init predictions api (must pass string, not Path)
+                    vals = [COCOeval(anno, pred, "bbox"), COCOeval(anno, pred, "segm")]
+                else:
+                    from lvis import LVIS, LVISEval
+
+                    anno = LVIS(str(anno_json))
+                    pred = anno._load_json(str(pred_json))
+                    vals = [LVISEval(anno, pred, "bbox"), LVISEval(anno, pred, "segm")]
+
+                for i, eval in enumerate(vals):
+                    eval.params.imgIds = [int(Path(x).stem) for x in self.dataloader.dataset.im_files]  # im to eval
                     eval.evaluate()
                     eval.accumulate()
                     eval.summarize()
+                    if self.is_lvis:
+                        eval.print_results()
                     idx = i * 4 + 2
-                    stats[self.metrics.keys[idx + 1]], stats[self.metrics.keys[idx]] = eval.stats[
-                        :2
-                    ]  # update mAP50-95 and mAP50
+                    # update mAP50-95 and mAP50
+                    stats[self.metrics.keys[idx + 1]], stats[self.metrics.keys[idx]] = (
+                        eval.stats[:2] if self.is_coco else [eval.results["AP"], eval.results["AP50"]]
+                    )
+                    if self.is_lvis:
+                        tag = "B" if i == 0 else "M"
+                        stats[f"metrics/APr({tag})"] = eval.results["APr"]
+                        stats[f"metrics/APc({tag})"] = eval.results["APc"]
+                        stats[f"metrics/APf({tag})"] = eval.results["APf"]
+
+                if self.is_lvis:
+                    stats["fitness"] = stats["metrics/mAP50-95(B)"]
+
             except Exception as e:
-                LOGGER.warning(f"pycocotools unable to run: {e}")
+                LOGGER.warning(f"{pkg} unable to run: {e}")
         return stats

ultralytics/models/yolo/yoloe/__init__.py

@@ -0,0 +1,21 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .predict import YOLOEVPDetectPredictor, YOLOEVPSegPredictor
+from .train import YOLOEPEFreeTrainer, YOLOEPETrainer, YOLOETrainer, YOLOEVPTrainer
+from .train_seg import YOLOEPESegTrainer, YOLOESegTrainer, YOLOESegTrainerFromScratch, YOLOESegVPTrainer
+from .val import YOLOEDetectValidator, YOLOESegValidator
+
+__all__ = [
+    "YOLOETrainer",
+    "YOLOEPETrainer",
+    "YOLOESegTrainer",
+    "YOLOEDetectValidator",
+    "YOLOESegValidator",
+    "YOLOEPESegTrainer",
+    "YOLOESegTrainerFromScratch",
+    "YOLOESegVPTrainer",
+    "YOLOEVPTrainer",
+    "YOLOEPEFreeTrainer",
+    "YOLOEVPDetectPredictor",
+    "YOLOEVPSegPredictor",
+]

ultralytics/models/yolo/yoloe/predict.py

@@ -0,0 +1,170 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+
+import numpy as np
+import torch
+
+from ultralytics.data.augment import LoadVisualPrompt
+from ultralytics.models.yolo.detect import DetectionPredictor
+from ultralytics.models.yolo.segment import SegmentationPredictor
+
+
+class YOLOEVPDetectPredictor(DetectionPredictor):
+    """
+    A mixin class for YOLO-EVP (Enhanced Visual Prompting) predictors.
+
+    This mixin provides common functionality for YOLO models that use visual prompting, including
+    model setup, prompt handling, and preprocessing transformations.
+
+    Attributes:
+        model (torch.nn.Module): The YOLO model for inference.
+        device (torch.device): Device to run the model on (CPU or CUDA).
+        prompts (dict): Visual prompts containing class indices and bounding boxes or masks.
+
+    Methods:
+        setup_model: Initialize the YOLO model and set it to evaluation mode.
+        set_return_vpe: Set whether to return visual prompt embeddings.
+        set_prompts: Set the visual prompts for the model.
+        pre_transform: Preprocess images and prompts before inference.
+        inference: Run inference with visual prompts.
+    """
+
+    def setup_model(self, model, verbose=True):
+        """
+        Sets up the model for prediction.
+
+        Args:
+            model (torch.nn.Module): Model to load or use.
+            verbose (bool): If True, provides detailed logging.
+        """
+        super().setup_model(model, verbose=verbose)
+        self.done_warmup = True
+
+    def set_prompts(self, prompts):
+        """
+        Set the visual prompts for the model.
+
+        Args:
+            prompts (dict): Dictionary containing class indices and bounding boxes or masks.
+                Must include a 'cls' key with class indices.
+        """
+        self.prompts = prompts
+
+    def pre_transform(self, im):
+        """
+        Preprocess images and prompts before inference.
+
+        This method applies letterboxing to the input image and transforms the visual prompts
+        (bounding boxes or masks) accordingly.
+
+        Args:
+            im (list): List containing a single input image.
+
+        Returns:
+            (list): Preprocessed image ready for model inference.
+
+        Raises:
+            ValueError: If neither valid bounding boxes nor masks are provided in the prompts.
+        """
+        img = super().pre_transform(im)
+        bboxes = self.prompts.pop("bboxes", None)
+        masks = self.prompts.pop("masks", None)
+        category = self.prompts["cls"]
+        if len(img) == 1:
+            visuals = self._process_single_image(img[0].shape[:2], im[0].shape[:2], category, bboxes, masks)
+            self.prompts = visuals.unsqueeze(0).to(self.device)  # (1, N, H, W)
+        else:
+            # NOTE: only supports bboxes as prompts for now
+            assert bboxes is not None, f"Expected bboxes, but got {bboxes}!"
+            # NOTE: needs List[np.ndarray]
+            assert isinstance(bboxes, list) and all(isinstance(b, np.ndarray) for b in bboxes), (
+                f"Expected List[np.ndarray], but got {bboxes}!"
+            )
+            assert isinstance(category, list) and all(isinstance(b, np.ndarray) for b in category), (
+                f"Expected List[np.ndarray], but got {category}!"
+            )
+            assert len(im) == len(category) == len(bboxes), (
+                f"Expected same length for all inputs, but got {len(im)}vs{len(category)}vs{len(bboxes)}!"
+            )
+            visuals = [
+                self._process_single_image(img[i].shape[:2], im[i].shape[:2], category[i], bboxes[i])
+                for i in range(len(img))
+            ]
+            self.prompts = torch.nn.utils.rnn.pad_sequence(visuals, batch_first=True).to(self.device)
+
+        return img
+
+    def _process_single_image(self, dst_shape, src_shape, category, bboxes=None, masks=None):
+        """
+        Processes a single image by resizing bounding boxes or masks and generating visuals.
+
+        Args:
+            dst_shape (tuple): The target shape (height, width) of the image.
+            src_shape (tuple): The original shape (height, width) of the image.
+            category (str): The category of the image for visual prompts.
+            bboxes (list | np.ndarray, optional): A list of bounding boxes in the format [x1, y1, x2, y2]. Defaults to None.
+            masks (np.ndarray, optional): A list of masks corresponding to the image. Defaults to None.
+
+        Returns:
+            visuals: The processed visuals for the image.
+
+        Raises:
+            ValueError: If neither `bboxes` nor `masks` are provided.
+        """
+        if bboxes is not None and len(bboxes):
+            bboxes = np.array(bboxes, dtype=np.float32)
+            if bboxes.ndim == 1:
+                bboxes = bboxes[None, :]
+            # Calculate scaling factor and adjust bounding boxes
+            gain = min(dst_shape[0] / src_shape[0], dst_shape[1] / src_shape[1])  # gain = old / new
+            bboxes *= gain
+            bboxes[..., 0::2] += round((dst_shape[1] - src_shape[1] * gain) / 2 - 0.1)
+            bboxes[..., 1::2] += round((dst_shape[0] - src_shape[0] * gain) / 2 - 0.1)
+        elif masks is not None:
+            # Resize and process masks
+            resized_masks = super().pre_transform(masks)
+            masks = np.stack(resized_masks)  # (N, H, W)
+            masks[masks == 114] = 0  # Reset padding values to 0
+        else:
+            raise ValueError("Please provide valid bboxes or masks")
+
+        # Generate visuals using the visual prompt loader
+        return LoadVisualPrompt().get_visuals(category, dst_shape, bboxes, masks)
+
+    def inference(self, im, *args, **kwargs):
+        """
+        Run inference with visual prompts.
+
+        Args:
+            im (torch.Tensor): Input image tensor.
+            *args (Any): Variable length argument list.
+            **kwargs (Any): Arbitrary keyword arguments.
+
+        Returns:
+            (torch.Tensor): Model prediction results.
+        """
+        return super().inference(im, vpe=self.prompts, *args, **kwargs)
+
+    def get_vpe(self, source):
+        """
+        Processes the source to get the visual prompt embeddings (VPE).
+
+        Args:
+            source (str | Path | int | PIL.Image | np.ndarray | torch.Tensor | List | Tuple): The source
+                of the image to make predictions on. Accepts various types including file paths, URLs, PIL
+                images, numpy arrays, and torch tensors.
+
+        Returns:
+            (torch.Tensor): The visual prompt embeddings (VPE) from the model.
+        """
+        self.setup_source(source)
+        assert len(self.dataset) == 1, "get_vpe only supports one image!"
+        for _, im0s, _ in self.dataset:
+            im = self.preprocess(im0s)
+            return self.model(im, vpe=self.prompts, return_vpe=True)
+
+
+class YOLOEVPSegPredictor(YOLOEVPDetectPredictor, SegmentationPredictor):
+    """Predictor for YOLOE VP segmentation."""
+
+    pass