dgenerate-ultralytics-headless 8.3.181__py3-none-any.whl → 8.3.183__py3-none-any.whl

{dgenerate_ultralytics_headless-8.3.181.dist-info → dgenerate_ultralytics_headless-8.3.183.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dgenerate-ultralytics-headless
-Version: 8.3.181
+Version: 8.3.183
 Summary: Automatically built Ultralytics package with python-opencv-headless dependency instead of python-opencv
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{dgenerate_ultralytics_headless-8.3.181.dist-info → dgenerate_ultralytics_headless-8.3.183.dist-info}/RECORD

@@ -1,14 +1,15 @@
-dgenerate_ultralytics_headless-8.3.181.dist-info/licenses/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
+dgenerate_ultralytics_headless-8.3.183.dist-info/licenses/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
 tests/__init__.py,sha256=b4KP5_q-2IO8Br8YHOSLYnn7IwZS81l_vfEF2YPa2lM,894
 tests/conftest.py,sha256=LXtQJcFNWPGuzauTGkiXgsvVC3llJKfg22WcmhRzuQc,2593
 tests/test_cli.py,sha256=EMf5gTAopOnIz8VvzaM-Qb044o7D0flnUHYQ-2ffOM4,5670
-tests/test_cuda.py,sha256=-nQsfF3lGfqLm6cIeu_BCiXqLj7HzpL7R1GzPEc6z2I,8128
+tests/test_cuda.py,sha256=7RAMC1DoXpsRvH0Jfyo9cqHkaJZWcWeqniCW5BW87hY,8228
 tests/test_engine.py,sha256=Jpt2KVrltrEgh2-3Ykouz-2Z_2fza0eymL5ectRXadM,4922
 tests/test_exports.py,sha256=CY-4xVZlVM16vdyIC0mSR3Ix59aiZm1qjFGIhSNmB20,11007
 tests/test_integrations.py,sha256=kl_AKmE_Qs1GB0_91iVwbzNxofm_hFTt0zzU6JF-pg4,6323
-tests/test_python.py,sha256=-qvdeg-hEcKU5mWSDEU24iFZ-i8FAwQRznSXpkp6WQ4,27928
+tests/test_python.py,sha256=JbOB6pbTkoQtPCjkl_idagV0_W2QLWGbsh2IvGmru0M,28274
 tests/test_solutions.py,sha256=tuf6n_fsI8KvSdJrnc-cqP2qYdiYqCWuVrx0z9dOz3Q,13213
-ultralytics/__init__.py,sha256=OqBNN1EOKn4_vq1OWj-ax36skQmjTn4HuE7IOLGpaI0,730
+ultralytics/__init__.py,sha256=hpM4gHWSSIAjKjCF3tuvHLzoZadJLeOZVINd-nMfpm8,730
+ultralytics/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
 ultralytics/assets/bus.jpg,sha256=wCAZxJecGR63Od3ZRERe9Aja1Weayrb9Ug751DS_vGM,137419
 ultralytics/assets/zidane.jpg,sha256=Ftc4aeMmen1O0A3o6GCDO9FlfBslLpTAw0gnetx7bts,50427
 ultralytics/cfg/__init__.py,sha256=Uj1br3-NVFvP6VY5CL4PK63mAQAom93XFC5cqSbM6t4,39887
@@ -109,7 +110,7 @@ ultralytics/data/annotator.py,sha256=uAgd7K-yudxiwdNqHz0ubfFg5JsfNlae4cgxdvCMyuY
 ultralytics/data/augment.py,sha256=Ps1s-ug_oXdyAz4Jyur6OmxzRlyzwP3VP-3hDalSxj8,132959
 ultralytics/data/base.py,sha256=mRcuehK1thNuuzQGL6D1AaZkod71oHRdYTod_zdQZQg,19688
 ultralytics/data/build.py,sha256=TfMLSPMbE2hGZVMLl178NTFrihC1-50jNOt1ex9elxw,11480
-ultralytics/data/converter.py,sha256=dExElV0vWd4EmDtZaFMC0clEmLdjRDIdFiXf01PUvQA,27134
+ultralytics/data/converter.py,sha256=1YJNJcOyzdnc-LmJw4oFbeJemHOWwx9z8m0V2QA_lzc,32053
 ultralytics/data/dataset.py,sha256=GhoFzBiuGvTr_5-3pzgWu6D_3aQVwW-hcS7kCo8XscM,36752
 ultralytics/data/loaders.py,sha256=u9sExTGPy1iiqVd_p29zVoEkQ3C36g2rE0FEbYPET0A,31767
 ultralytics/data/split.py,sha256=F6O73bAbESj70FQZzqkydXQeXgPXGHGiC06b5MkLHjQ,5109
@@ -124,7 +125,7 @@ ultralytics/engine/exporter.py,sha256=Vr7K8Yf3wyf91ZvDpRosAohwa_W0oe4qW-JvqigCPf
 ultralytics/engine/model.py,sha256=877u2n0ISz2COOYtEMUqQe0E-HHB4Atb2DuH1XCE98k,53530
 ultralytics/engine/predictor.py,sha256=iXnUB-tvBHtVpKbB-5EKs1wSREBIerdUxWx39MaFYuk,22485
 ultralytics/engine/results.py,sha256=QcHcbPVlLBiy_APwABr-T5K65HR8Bl1rRzxawjjP76E,71873
-ultralytics/engine/trainer.py,sha256=28FeqASvQRxCaK96SXDM-BfPJjqy5KNiWhf8v6GXTug,39785
+ultralytics/engine/trainer.py,sha256=rOzMFt7b5PzVJZ3cOl80wxQAKUTP7XhnMHfXSaVRbOE,40193
 ultralytics/engine/tuner.py,sha256=sfQ8_yzgLNcGlKyz9b2vAzyggGZXiQzdZ5tKstyqjHM,12825
 ultralytics/engine/validator.py,sha256=g0StH6WOn95zBN-hULDAR5Uug1pU2YkaeNH3zzq3SVg,16573
 ultralytics/hub/__init__.py,sha256=ulPtceI3hqud03mvqoXccBaa1e4nveYwC9cddyuBUlo,6599
@@ -151,13 +152,13 @@ ultralytics/models/sam/__init__.py,sha256=iR7B06rAEni21eptg8n4rLOP0Z_qV9y9PL-L93
 ultralytics/models/sam/amg.py,sha256=IpcuIfC5KBRiF4sdrsPl1ecWEJy75axo1yG23r5BFsw,11783
 ultralytics/models/sam/build.py,sha256=J6n-_QOYLa63jldEZmhRe9D3Is_AJE8xyZLUjzfRyTY,12629
 ultralytics/models/sam/model.py,sha256=j1TwsLmtxhiXyceU31VPzGVkjRXGylphKrdPSzUJRJc,7231
-ultralytics/models/sam/predict.py,sha256=awE_46I-GmYRIeDDLmGIdaYwJvPeSbw316DyanrA1Ys,86453
+ultralytics/models/sam/predict.py,sha256=R32JjExRBL5c2zBcDdauhX4UM8E8kMrBLoa0sZ9vk6I,86494
 ultralytics/models/sam/modules/__init__.py,sha256=lm6MckFYCPTbqIoX7w0s_daxdjNeBeKW6DXppv1-QUM,70
 ultralytics/models/sam/modules/blocks.py,sha256=lnMhnexvXejzhixWRQQyqjrpALoIhuOSwnSGW-c9kZk,46089
 ultralytics/models/sam/modules/decoders.py,sha256=U9jqFRkD0JmO3eugSmwLD0sQkiGqJJLympWNO83osGM,25638
 ultralytics/models/sam/modules/encoders.py,sha256=srtxrfy3SfUarkC41L1S8tY4GdFueUuR2qQDFZ6ZPl4,37362
 ultralytics/models/sam/modules/memory_attention.py,sha256=F1XJAxSwho2-LMlrao_ij0MoALTvhkK-OVghi0D4cU0,13651
-ultralytics/models/sam/modules/sam.py,sha256=ACI2wA-FiWwj5ctHMHJIi_ZMw4ujrBkHEaZ77X1De_Y,55649
+ultralytics/models/sam/modules/sam.py,sha256=CjM4M2PfRltQFnHFOp2G6QAdYk9BxWlurx82FSX_TYo,55760
 ultralytics/models/sam/modules/tiny_encoder.py,sha256=lmUIeZ9-3M-C3YmJBs13W6t__dzeJloOl0qFR9Ll8ew,42241
 ultralytics/models/sam/modules/transformer.py,sha256=xc2g6gb0jvr7cJkHkzIbZOGcTrmsOn2ojvuH-MVIMVs,14953
 ultralytics/models/sam/modules/utils.py,sha256=-PYSLExtBajbotBdLan9J07aFaeXJ03WzopAv4JcYd4,16022
@@ -193,7 +194,7 @@ ultralytics/models/yolo/yoloe/__init__.py,sha256=6SLytdJtwu37qewf7CobG7C7Wl1m-xt
 ultralytics/models/yolo/yoloe/predict.py,sha256=GmQxCQe7sLomAujde53jQzquzryNn6fEjS4Oalf3mPs,7124
 ultralytics/models/yolo/yoloe/train.py,sha256=XYpQYSnSD8vi_9VSj_S5oIsNUEqm3e66vPT8rNFI_HY,14086
 ultralytics/models/yolo/yoloe/train_seg.py,sha256=aCV7M8oQOvODFnU4piZdJh3tIrBJYAzZfRVRx1vRgxo,4956
-ultralytics/models/yolo/yoloe/val.py,sha256=yebPkxwKKt__cY05Zbh1YXg4_BKzzpcDc3Cv3FJ5SAA,9769
+ultralytics/models/yolo/yoloe/val.py,sha256=2NuERI3B3WeED658Cat1xL2SVpORUHlCHCWI3L8pJJc,9784
 ultralytics/nn/__init__.py,sha256=rjociYD9lo_K-d-1s6TbdWklPLjTcEHk7OIlRDJstIE,615
 ultralytics/nn/autobackend.py,sha256=UM9ObXeLB0lgak1Q5oSi2IA-R_Owr6NdJNBAsA3mSbo,41790
 ultralytics/nn/tasks.py,sha256=vw_TNacAv-RN24rusFzKuYL6qRBD7cve8EpB7gOlU_8,72505
@@ -250,7 +251,7 @@ ultralytics/utils/loss.py,sha256=fbOWc3Iu0QOJiWbi-mXWA9-1otTYlehtmUsI7os7ydM,397
 ultralytics/utils/metrics.py,sha256=tQjYxPd0dSzjucVyI1evIISunyYRkABXMXVQo64mAUE,68756
 ultralytics/utils/ops.py,sha256=8d60fbpntrexK3gPoLUS6mWAYGrtrQaQCOYyRJsCjuI,34521
 ultralytics/utils/patches.py,sha256=PPWiKzwGbCvuawLzDKVR8tWOQAlZbJBi8g_-A6eTCYA,6536
-ultralytics/utils/plotting.py,sha256=IEugKlTITLxArZjbSr7i_cTaHHAqNwVVk08Ak7I_ZdM,47169
+ultralytics/utils/plotting.py,sha256=4TG_J8rz9VVPrOXbdjRHPJZVgJrFYVmEYE0BcVDdolc,47745
 ultralytics/utils/tal.py,sha256=aXawOnhn8ni65tJWIW-PYqWr_TRvltbHBjrTo7o6lDQ,20924
 ultralytics/utils/torch_utils.py,sha256=D76Pvmw5OKh-vd4aJkOMO0dSLbM5WzGr7Hmds54hPEk,39233
 ultralytics/utils/triton.py,sha256=M7qe4RztiADBJQEWQKaIQsp94ERFJ_8_DUHDR6TXEOM,5410
@@ -266,8 +267,8 @@ ultralytics/utils/callbacks/neptune.py,sha256=j8pecmlcsM8FGzLKWoBw5xUsi5t8E5HuxY
 ultralytics/utils/callbacks/raytune.py,sha256=S6Bq16oQDQ8BQgnZzA0zJHGN_BBr8iAM_WtGoLiEcwg,1283
 ultralytics/utils/callbacks/tensorboard.py,sha256=MDPBW7aDes-66OE6YqKXXvqA_EocjzEMHWGM-8z9vUQ,5281
 ultralytics/utils/callbacks/wb.py,sha256=Tm_-aRr2CN32MJkY9tylpMBJkb007-MSRNSQ7rDJ5QU,7521
-dgenerate_ultralytics_headless-8.3.181.dist-info/METADATA,sha256=6a7UOAonIPqJS7OoY1QQ6pBR1hIhPk4Tu5Rb-RSlINU,38727
-dgenerate_ultralytics_headless-8.3.181.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dgenerate_ultralytics_headless-8.3.181.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
-dgenerate_ultralytics_headless-8.3.181.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
-dgenerate_ultralytics_headless-8.3.181.dist-info/RECORD,,
+dgenerate_ultralytics_headless-8.3.183.dist-info/METADATA,sha256=u9Pxkr9kFuSZDzhiV3hvbkFV8Br-rfRoWSABCSUcQ1E,38727
+dgenerate_ultralytics_headless-8.3.183.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dgenerate_ultralytics_headless-8.3.183.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
+dgenerate_ultralytics_headless-8.3.183.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
+dgenerate_ultralytics_headless-8.3.183.dist-info/RECORD,,

tests/test_cuda.py

@@ -67,7 +67,7 @@ def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
         half=half,
         batch=batch,
         simplify=simplify,
-        nms=nms,
+        nms=nms and task != "obb",  # disable NMS for OBB task for now on T4 instance
         device=DEVICES[0],
     )
     YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32, device=DEVICES[0])  # exported model inference
@@ -163,7 +163,7 @@ def test_autobatch():
 
 
 @pytest.mark.slow
-@pytest.mark.skipif(not DEVICES, reason="No CUDA devices available")
+@pytest.mark.skipif(True, reason="Skip for now since T4 instance does not support TensorRT > 10.0")
 def test_utils_benchmarks():
     """Profile YOLO models for performance benchmarks."""
     from ultralytics.utils.benchmarks import ProfileModels

tests/test_python.py

@@ -228,6 +228,15 @@ def test_train_scratch():
     model(SOURCE)
 
 
+@pytest.mark.skipif(not ONLINE, reason="environment is offline")
+def test_train_ndjson():
+    """Test training the YOLO model using NDJSON format dataset."""
+    model = YOLO(WEIGHTS_DIR / "yolo11n.pt")
+    model.train(
+        data="https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8-ndjson.ndjson", epochs=1, imgsz=32
+    )
+
+
 @pytest.mark.parametrize("scls", [False, True])
 def test_train_pretrained(scls):
     """Test training of the YOLO model starting from a pre-trained checkpoint."""

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.181"
+__version__ = "8.3.183"
 
 import os
 

ultralytics/data/converter.py

@@ -1,5 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+import asyncio
 import json
 import random
 import shutil
@@ -12,7 +13,8 @@ import cv2
 import numpy as np
 from PIL import Image
 
-from ultralytics.utils import DATASETS_DIR, LOGGER, NUM_THREADS, TQDM
+from ultralytics.utils import DATASETS_DIR, LOGGER, NUM_THREADS, TQDM, YAML
+from ultralytics.utils.checks import check_file, check_requirements
 from ultralytics.utils.downloads import download, zip_directory
 from ultralytics.utils.files import increment_path
 
@@ -754,3 +756,112 @@ def convert_to_multispectral(path: Union[str, Path], n_channels: int = 10, repla
         multispectral = f(target_wavelengths)
         cv2.imwritemulti(str(output_path), np.clip(multispectral, 0, 255).astype(np.uint8).transpose(2, 0, 1))
         LOGGER.info(f"Converted {output_path}")
+
+
+async def convert_ndjson_to_yolo(ndjson_path: Union[str, Path], output_path: Optional[Union[str, Path]] = None) -> Path:
+    """
+    Convert NDJSON dataset format to Ultralytics YOLO11 dataset structure.
+
+    This function converts datasets stored in NDJSON (Newline Delimited JSON) format to the standard YOLO
+    format with separate directories for images and labels. It supports parallel processing for efficient
+    conversion of large datasets and can download images from URLs if they don't exist locally.
+
+    The NDJSON format consists of:
+    - First line: Dataset metadata with class names and configuration
+    - Subsequent lines: Individual image records with annotations and optional URLs
+
+    Args:
+        ndjson_path (Union[str, Path]): Path to the input NDJSON file containing dataset information.
+        output_path (Optional[Union[str, Path]], optional): Directory where the converted YOLO dataset
+            will be saved. If None, uses the parent directory of the NDJSON file. Defaults to None.
+
+    Returns:
+        (Path): Path to the generated data.yaml file that can be used for YOLO training.
+
+    Examples:
+        Convert a local NDJSON file:
+        >>> yaml_path = convert_ndjson_to_yolo("dataset.ndjson")
+        >>> print(f"Dataset converted to: {yaml_path}")
+
+        Convert with custom output directory:
+        >>> yaml_path = convert_ndjson_to_yolo("dataset.ndjson", output_path="./converted_datasets")
+
+        # Use with YOLO training
+        >>> from ultralytics import YOLO
+        >>> model = YOLO("yolo11n.pt")
+        >>> model.train(data="https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8-ndjson.ndjson")
+    """
+    check_requirements("aiohttp")
+    import aiohttp
+
+    ndjson_path = Path(check_file(ndjson_path))
+    output_path = Path(output_path or DATASETS_DIR)
+    with open(ndjson_path) as f:
+        lines = [json.loads(line.strip()) for line in f if line.strip()]
+
+    dataset_record, image_records = lines[0], lines[1:]
+    dataset_dir = output_path / ndjson_path.stem
+    splits = {record["split"] for record in image_records}
+
+    # Create directories and prepare YAML structure
+    dataset_dir.mkdir(parents=True, exist_ok=True)
+    data_yaml = dict(dataset_record)
+    data_yaml["names"] = {int(k): v for k, v in dataset_record.get("class_names", {}).items()}
+
+    for split in sorted(splits):
+        (dataset_dir / "images" / split).mkdir(parents=True, exist_ok=True)
+        (dataset_dir / "labels" / split).mkdir(parents=True, exist_ok=True)
+        data_yaml[split] = f"images/{split}"
+
+    async def process_record(session, semaphore, record):
+        """Process single image record with async session."""
+        async with semaphore:
+            split, original_name = record["split"], record["file"]
+            label_path = dataset_dir / "labels" / split / f"{Path(original_name).stem}.txt"
+            image_path = dataset_dir / "images" / split / original_name
+
+            annotations = record.get("annotations", {})
+            lines_to_write = []
+            for key in annotations.keys():
+                lines_to_write = [" ".join(map(str, item)) for item in annotations[key]]
+                break
+            if "classification" in annotations:
+                lines_to_write = [str(cls) for cls in annotations["classification"]]
+
+            label_path.write_text("\n".join(lines_to_write) + "\n" if lines_to_write else "")
+
+            if http_url := record.get("url"):
+                if not image_path.exists():
+                    try:
+                        async with session.get(http_url, timeout=aiohttp.ClientTimeout(total=30)) as response:
+                            response.raise_for_status()
+                            with open(image_path, "wb") as f:
+                                async for chunk in response.content.iter_chunked(8192):
+                                    f.write(chunk)
+                        return True
+                    except Exception as e:
+                        LOGGER.warning(f"Failed to download {http_url}: {e}")
+                        return False
+            return True
+
+    # Process all images with async downloads
+    semaphore = asyncio.Semaphore(64)
+    async with aiohttp.ClientSession() as session:
+        pbar = TQDM(
+            total=len(image_records),
+            desc=f"Converting {ndjson_path.name} → {dataset_dir} ({len(image_records)} images)",
+        )
+
+        async def tracked_process(record):
+            result = await process_record(session, semaphore, record)
+            pbar.update(1)
+            return result
+
+        await asyncio.gather(*[tracked_process(record) for record in image_records])
+        pbar.close()
+
+    # Write data.yaml
+    yaml_path = dataset_dir / "data.yaml"
+    YAML.save(yaml_path, data_yaml)
+
+    return yaml_path

ultralytics/engine/trainer.py

@@ -598,6 +598,15 @@ class BaseTrainer:
         try:
             if self.args.task == "classify":
                 data = check_cls_dataset(self.args.data)
+            elif self.args.data.rsplit(".", 1)[-1] == "ndjson":
+                # Convert NDJSON to YOLO format
+                import asyncio
+
+                from ultralytics.data.converter import convert_ndjson_to_yolo
+
+                yaml_path = asyncio.run(convert_ndjson_to_yolo(self.args.data))
+                self.args.data = str(yaml_path)
+                data = check_det_dataset(self.args.data)
             elif self.args.data.rsplit(".", 1)[-1] in {"yaml", "yml"} or self.args.task in {
                 "detect",
                 "segment",

ultralytics/models/sam/modules/sam.py

@@ -1032,3 +1032,4 @@ class SAM2Model(torch.nn.Module):
         self.image_size = imgsz[0]
         self.sam_prompt_encoder.input_image_size = imgsz
         self.sam_prompt_encoder.image_embedding_size = [x // 16 for x in imgsz]  # fixed ViT patch size of 16
+        self.sam_image_embedding_size = self.image_size // self.backbone_stride  # update image embedding size

ultralytics/models/sam/predict.py

@@ -1283,6 +1283,7 @@ class SAM2VideoPredictor(SAM2Predictor):
             - If `batch` is greater than 1, the features are expanded to fit the batch size.
             - The method leverages the model's `_prepare_backbone_features` method to prepare the backbone features.
         """
+        self.model.set_imgsz(self.imgsz)
         backbone_out = self.model.forward_image(im)
         if batch > 1:  # expand features if there's more than one prompt
             for i, feat in enumerate(backbone_out["backbone_fpn"]):

ultralytics/models/yolo/yoloe/val.py

@@ -181,7 +181,7 @@ class YOLOEDetectValidator(DetectionValidator):
         else:
             if refer_data is not None:
                 assert load_vp, "Refer data is only used for visual prompt validation."
-            self.device = select_device(self.args.device)
+            self.device = select_device(self.args.device, verbose=False)
 
             if isinstance(model, (str, Path)):
                 from ultralytics.nn.tasks import attempt_load_weights

ultralytics/py.typed

@@ -0,0 +1 @@
+partial

ultralytics/utils/plotting.py

@@ -361,38 +361,46 @@ class Annotator:
                     lineType=cv2.LINE_AA,
                 )
 
-    def masks(self, masks, colors, im_gpu, alpha: float = 0.5, retina_masks: bool = False):
+    def masks(self, masks, colors, im_gpu: torch.Tensor = None, alpha: float = 0.5, retina_masks: bool = False):
         """
         Plot masks on image.
 
         Args:
-            masks (torch.Tensor): Predicted masks on cuda, shape: [n, h, w]
+            masks (torch.Tensor | np.ndarray): Predicted masks with shape: [n, h, w]
             colors (List[List[int]]): Colors for predicted masks, [[r, g, b] * n]
-            im_gpu (torch.Tensor): Image is in cuda, shape: [3, h, w], range: [0, 1]
+            im_gpu (torch.Tensor | None): Image is in cuda, shape: [3, h, w], range: [0, 1]
             alpha (float, optional): Mask transparency: 0.0 fully transparent, 1.0 opaque.
             retina_masks (bool, optional): Whether to use high resolution masks or not.
         """
         if self.pil:
             # Convert to numpy first
             self.im = np.asarray(self.im).copy()
-        if len(masks) == 0:
-            self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
-        if im_gpu.device != masks.device:
-            im_gpu = im_gpu.to(masks.device)
-        colors = torch.tensor(colors, device=masks.device, dtype=torch.float32) / 255.0  # shape(n,3)
-        colors = colors[:, None, None]  # shape(n,1,1,3)
-        masks = masks.unsqueeze(3)  # shape(n,h,w,1)
-        masks_color = masks * (colors * alpha)  # shape(n,h,w,3)
-
-        inv_alpha_masks = (1 - masks * alpha).cumprod(0)  # shape(n,h,w,1)
-        mcs = masks_color.max(dim=0).values  # shape(n,h,w,3)
-
-        im_gpu = im_gpu.flip(dims=[0])  # flip channel
-        im_gpu = im_gpu.permute(1, 2, 0).contiguous()  # shape(h,w,3)
-        im_gpu = im_gpu * inv_alpha_masks[-1] + mcs
-        im_mask = im_gpu * 255
-        im_mask_np = im_mask.byte().cpu().numpy()
-        self.im[:] = im_mask_np if retina_masks else ops.scale_image(im_mask_np, self.im.shape)
+        if im_gpu is None:
+            assert isinstance(masks, np.ndarray), "`masks` must be a np.ndarray if `im_gpu` is not provided."
+            overlay = self.im.copy()
+            for i, mask in enumerate(masks):
+                overlay[mask.astype(bool)] = colors[i]
+            self.im = cv2.addWeighted(self.im, 1 - alpha, overlay, alpha, 0)
+        else:
+            assert isinstance(masks, torch.Tensor), "`masks` must be a torch.Tensor if `im_gpu` is provided."
+            if len(masks) == 0:
+                self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
+            if im_gpu.device != masks.device:
+                im_gpu = im_gpu.to(masks.device)
+            colors = torch.tensor(colors, device=masks.device, dtype=torch.float32) / 255.0  # shape(n,3)
+            colors = colors[:, None, None]  # shape(n,1,1,3)
+            masks = masks.unsqueeze(3)  # shape(n,h,w,1)
+            masks_color = masks * (colors * alpha)  # shape(n,h,w,3)
+
+            inv_alpha_masks = (1 - masks * alpha).cumprod(0)  # shape(n,h,w,1)
+            mcs = masks_color.max(dim=0).values  # shape(n,h,w,3)
+
+            im_gpu = im_gpu.flip(dims=[0])  # flip channel
+            im_gpu = im_gpu.permute(1, 2, 0).contiguous()  # shape(h,w,3)
+            im_gpu = im_gpu * inv_alpha_masks[-1] + mcs
+            im_mask = im_gpu * 255
+            im_mask_np = im_mask.byte().cpu().numpy()
+            self.im[:] = im_mask_np if retina_masks else ops.scale_image(im_mask_np, self.im.shape)
         if self.pil:
             # Convert im back to PIL and update draw
             self.fromarray(self.im)