Visualize Ultralytics Yolo models

examples/computer_vision/.gitignore

@@ -0,0 +1,3 @@
+.datachain
+output
+*.pt

examples/computer_vision/ultralytics-bbox.py

@@ -1,25 +1,37 @@
 import os
-
-os.environ["YOLO_VERBOSE"] = "false"
-
-
 from io import BytesIO
 
+from numpy import asarray
 from PIL import Image
 from ultralytics import YOLO
 
-from datachain import C, DataChain, File
+from datachain import DataChain, File
 from datachain.model.ultralytics import YoloBBoxes
+from datachain.toolkit.ultralytics import visualize_yolo
+
+OUTPUT_DIR = "output/bbox"
 
 
 def process_bboxes(yolo: YOLO, file: File) -> YoloBBoxes:
-    results = yolo(Image.open(BytesIO(file.read())))
-    return YoloBBoxes.from_results(results)
+    # read image
+    img = Image.open(BytesIO(file.read()))
+
+    # detect objects using YOLO model
+    results = yolo(img, verbose=False)
+    # convert results to YoloBBoxes signal
+    signal = YoloBBoxes.from_results(results)
+
+    # visualize results
+    img2 = visualize_yolo(asarray(img), signal)
+    img2.save(f"{OUTPUT_DIR}/{file.get_file_stem()}.jpg")
+
+    return signal
+
 
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 (
-    DataChain.from_storage("gs://datachain-demo/openimages-v6-test-jsonpairs/")
-    .filter(C("file.path").glob("*.jpg"))
+    DataChain.from_storage("gs://datachain-demo/coco2017/images")
     .limit(20)
     .setup(yolo=lambda: YOLO("yolo11n.pt"))
     .map(boxes=process_bboxes)

examples/computer_vision/ultralytics-pose.py

@@ -1,25 +1,37 @@
 import os
-
-os.environ["YOLO_VERBOSE"] = "false"
-
-
 from io import BytesIO
 
+from numpy import asarray
 from PIL import Image
 from ultralytics import YOLO
 
-from datachain import C, DataChain, File
+from datachain import DataChain, File
 from datachain.model.ultralytics import YoloPoses
+from datachain.toolkit.ultralytics import visualize_yolo
+
+OUTPUT_DIR = "output/pose"
 
 
 def process_poses(yolo: YOLO, file: File) -> YoloPoses:
-    results = yolo(Image.open(BytesIO(file.read())))
-    return YoloPoses.from_results(results)
+    # read image
+    img = Image.open(BytesIO(file.read()))
+
+    # detect objects using YOLO model
+    results = yolo(img, verbose=False)
+    # convert results to YoloPoses signal
+    signal = YoloPoses.from_results(results)
+
+    # visualize results
+    img2 = visualize_yolo(asarray(img), signal)
+    img2.save(f"{OUTPUT_DIR}/{file.get_file_stem()}.jpg")
+
+    return signal
+
 
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 (
-    DataChain.from_storage("gs://datachain-demo/openimages-v6-test-jsonpairs/")
-    .filter(C("file.path").glob("*.jpg"))
+    DataChain.from_storage("gs://datachain-demo/coco2017/images")
     .limit(20)
     .setup(yolo=lambda: YOLO("yolo11n-pose.pt"))
     .map(poses=process_poses)

examples/computer_vision/ultralytics-segment.py

@@ -1,25 +1,37 @@
 import os
-
-os.environ["YOLO_VERBOSE"] = "false"
-
-
 from io import BytesIO
 
+from numpy import asarray
 from PIL import Image
 from ultralytics import YOLO
 
-from datachain import C, DataChain, File
+from datachain import DataChain, File
 from datachain.model.ultralytics import YoloSegments
+from datachain.toolkit.ultralytics import visualize_yolo
+
+OUTPUT_DIR = "output/segment"
 
 
 def process_segments(yolo: YOLO, file: File) -> YoloSegments:
-    results = yolo(Image.open(BytesIO(file.read())))
-    return YoloSegments.from_results(results)
+    # read image
+    img = Image.open(BytesIO(file.read()))
+
+    # detect objects using YOLO model
+    results = yolo(img, verbose=False)
+    # convert results to YoloSegments signal
+    signal = YoloSegments.from_results(results)
+
+    # visualize results
+    img2 = visualize_yolo(asarray(img), signal)
+    img2.save(f"{OUTPUT_DIR}/{file.get_file_stem()}.jpg")
+
+    return signal
+
 
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 (
-    DataChain.from_storage("gs://datachain-demo/openimages-v6-test-jsonpairs/")
-    .filter(C("file.path").glob("*.jpg"))
+    DataChain.from_storage("gs://datachain-demo/coco2017/images")
     .limit(20)
     .setup(yolo=lambda: YOLO("yolo11n-seg.pt"))
     .map(segments=process_segments)

src/datachain/toolkit/ultralytics.py

@@ -0,0 +1,159 @@
+from typing import Union
+
+import numpy as np
+import torch
+from PIL import Image
+from ultralytics.data.utils import polygon2mask
+from ultralytics.engine.results import Results
+
+from datachain.model.ultralytics.bbox import YoloBBox, YoloBBoxes
+from datachain.model.ultralytics.pose import YoloPose, YoloPoses
+from datachain.model.ultralytics.segment import YoloSegment, YoloSegments
+
+YoloSignal = Union[YoloBBox, YoloBBoxes, YoloPose, YoloPoses, YoloSegment, YoloSegments]
+
+
+def _signal_to_results(img: np.ndarray, signal: YoloSignal) -> Results:
+    """Convert a YOLO signal to Ultralytics Results."""
+    # Convert RGB to BGR
+    if img.ndim == 3 and img.shape[2] == 3:
+        bgr_array = img[:, :, ::-1]
+    else:
+        # If the image is not RGB (e.g., grayscale or RGBA), use as is
+        bgr_array = img
+
+    names = {}
+    boxes_list = []
+    keypoints_list = []
+    masks_list = []
+
+    # Get the boxes, keypoints, and masks from the signal
+    if isinstance(signal, YoloBBox):
+        names[signal.cls] = signal.name
+        boxes_list.append(
+            torch.tensor([[*signal.box.coords, signal.confidence, signal.cls]])
+        )
+    elif isinstance(signal, YoloBBoxes):
+        for i, _ in enumerate(signal.cls):
+            names[signal.cls[i]] = signal.name[i]
+            boxes_list.append(
+                torch.tensor(
+                    [[*signal.box[i].coords, signal.confidence[i], signal.cls[i]]]
+                )
+            )
+    elif isinstance(signal, YoloPose):
+        names[signal.cls] = signal.name
+        boxes_list.append(
+            torch.tensor([[*signal.box.coords, signal.confidence, signal.cls]])
+        )
+        keypoints_list.append(
+            torch.tensor([list(zip(signal.pose.x, signal.pose.y, signal.pose.visible))])
+        )
+    elif isinstance(signal, YoloPoses):
+        for i, _ in enumerate(signal.cls):
+            names[signal.cls[i]] = signal.name[i]
+            boxes_list.append(
+                torch.tensor(
+                    [[*signal.box[i].coords, signal.confidence[i], signal.cls[i]]]
+                )
+            )
+            keypoints_list.append(
+                torch.tensor(
+                    [
+                        list(
+                            zip(
+                                signal.pose[i].x,
+                                signal.pose[i].y,
+                                signal.pose[i].visible,
+                            )
+                        )
+                    ]
+                )
+            )
+    elif isinstance(signal, YoloSegment):
+        names[signal.cls] = signal.name
+        boxes_list.append(
+            torch.tensor([[*signal.box.coords, signal.confidence, signal.cls]])
+        )
+        masks_list.append(
+            torch.tensor(
+                polygon2mask(
+                    img.shape[:2],
+                    [np.asarray(list(zip(signal.segment.x, signal.segment.y)))],
+                )
+            )
+        )
+    elif isinstance(signal, YoloSegments):
+        for i, _ in enumerate(signal.cls):
+            names[signal.cls[i]] = signal.name[i]
+            boxes_list.append(
+                torch.tensor(
+                    [[*signal.box[i].coords, signal.confidence[i], signal.cls[i]]]
+                )
+            )
+            masks_list.append(
+                torch.tensor(
+                    polygon2mask(
+                        img.shape[:2],
+                        [
+                            np.asarray(
+                                list(zip(signal.segment[i].x, signal.segment[i].y))
+                            )
+                        ],
+                    )
+                )
+            )
+
+    boxes = torch.cat(boxes_list, dim=0) if len(boxes_list) > 0 else None
+    keypoints = torch.cat(keypoints_list, dim=0) if len(keypoints_list) > 0 else None
+    masks = torch.stack(masks_list) if len(masks_list) > 0 else None
+
+    return Results(
+        bgr_array,
+        path="",
+        names=names,
+        boxes=boxes,
+        keypoints=keypoints,
+        masks=masks,
+    )
+
+
+def visualize_yolo(
+    img: np.ndarray,
+    signal: YoloSignal,
+    scale: float = 1.0,
+    line_width: int = 1,
+    font_size: int = 20,
+    kpt_radius: int = 3,
+) -> Image.Image:
+    """
+    Visualize signals detected by YOLO.
+
+    Args:
+        image (ndarray): The image to visualize as a NumPy array.
+        signal: The signal detected by YOLO. Possible signals are YoloBBox, YoloBBoxes,
+                YoloPose, YoloPoses, YoloSegment, and YoloSegments.
+        scale (float): The scale factor for the image. Default is 1.0.
+        line_width (int): The line width for drawing boxes and lines. Default is 1.
+        font_size (int): The font size for text. Default is 20.
+        kpt_radius (int): The radius for drawing keypoints. Default is 3.
+
+    Returns:
+        PIL.Image.Image: The image with the detected signals visualized.
+    """
+    results = _signal_to_results(img, signal)
+
+    im_bgr = results.plot(
+        line_width=line_width,
+        font_size=font_size,
+        kpt_radius=kpt_radius,
+    )
+
+    im_rgb = Image.fromarray(im_bgr[..., ::-1])
+
+    if scale != 1.0:
+        orig_height, orig_width = results.orig_shape
+        new_size = (int(orig_width * scale), int(orig_height * scale))
+        im_rgb = im_rgb.resize(new_size, Image.Resampling.LANCZOS)
+
+    return im_rgb