add yolo node

.gitignore

@@ -0,0 +1,5 @@
+.git
+.vscode
+**/__pycache__
+**/*.ckpt
+.DS_Store

README.md

@@ -0,0 +1,22 @@
+# ComfyUI-TrackingNodes
+ComfyUI nodes to track objects in videos using YOLOv8
+
+
+## Table of Contents
+- [Installation](#installation)
+  - [How to Install](#how-to-install)
+  - [How to Configure Models](#how-to-configure-models)
+- [Examples](#examples)
+
+## Installation
+
+### How to Install
+Clone or download this repo into your `ComfyUI/custom_nodes/` directory.
+Then run `pip install -r requirments.txt` within the cloned repo.
+
+### How to Configure Models
+Download one of the YOLOv8 models from https://github.com/ultralytics/ultralytics and place it into `ComfyUI/custom_nodes/ComfyUI-TrackingNodes/models/`
+
+## Accompanying Node Repos
+* [InstanceDiffusion](https://github.com/logtd/ComfyUI-InstanceDiffusion)
+* [AnimateDiff-Evolved](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved)

__init__.py

@@ -0,0 +1,10 @@
+from .nodes.yolo_tracker_node import YOLOTrackerNode
+
+
+NODE_CLASS_MAPPINGS = {
+    "YOLOTrackerNode": YOLOTrackerNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "YOLOTrackerNode": "YOLO BoundingBox Tracker",
+}

nodes/yolo_tracker_node.py

@@ -0,0 +1,95 @@
+import os
+import folder_paths
+import numpy as np
+import torch
+from PIL import Image
+import supervision as sv
+from ultralytics import YOLO
+
+folder_paths.folder_names_and_paths["yolov8"] = ([os.path.join(
+    folder_paths.models_dir, "yolov8")], folder_paths.supported_pt_extensions)
+
+
+def get_model_dir():
+    root_path = folder_paths.folder_names_and_paths['custom_nodes'][0][0]
+    path = os.path.join(root_path, 'ComfyUI-TrackingNodes', 'models')
+    return path
+
+
+def get_model_list() -> list[str]:
+    path = get_model_dir()
+    return os.listdir(path)
+
+
+class YOLOTrackerNode:
+    def __init__(self) -> None:
+        ...
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "model_name": (get_model_list(), ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "TRACKING")
+    FUNCTION = "track"
+    CATEGORY = "tracking"
+
+    def track(self, images, model_name):
+        B, H, W, C = images.shape
+        # load a custom model
+        model = YOLO(
+            f'{os.path.join(folder_paths.models_dir, "yolov8")}/{model_name}')
+        tracker = sv.ByteTrack()
+        box_annotator = sv.BoundingBoxAnnotator()
+        label_annotator = sv.LabelAnnotator()
+        annotated = []
+
+        tracked = {}
+
+        class_tracker_id_map = {}
+
+        for idx, image in enumerate(images):
+            image = (image.cpu().numpy() * 255).astype(np.uint8)
+            # bug with opencv2 and need to convert to pillow first
+            results = model(Image.fromarray(image))[0]
+            detections = sv.Detections.from_ultralytics(results)
+            detections = tracker.update_with_detections(detections)
+
+            for i in range(len(detections)):
+                detection = detections[i]
+                class_id = model.names[detection.class_id[0]].lower()
+                tracker_id = detection.tracker_id[0]
+                if class_id not in class_tracker_id_map:
+                    class_tracker_id_map[class_id] = {}
+
+                if tracker_id not in class_tracker_id_map[class_id]:
+                    next_id = len(class_tracker_id_map[class_id])
+                    class_tracker_id_map[class_id][tracker_id] = next_id
+
+                if class_id not in tracked:
+                    tracked[class_id] = {}
+                if tracker_id not in tracked[class_id]:
+                    tracked[class_id][tracker_id] = [None] * len(images)
+
+                tracked[class_id][tracker_id][idx] = list(
+                    map(lambda x: int(x), detection.xyxy[0])) + [W, H]
+
+            labels = [
+                f"#{tracker_id}.{results.names[class_id].lower()}"
+                for class_id, tracker_id
+                in zip(detections.class_id, detections.tracker_id)
+            ]
+
+            annotated_frame = box_annotator.annotate(
+                image.copy(), detections=detections)
+
+            annotated.append(torch.FloatTensor(label_annotator.annotate(
+                annotated_frame, detections=detections, labels=labels))/255.)
+
+        annotated = torch.stack(annotated)
+        del model
+        return (annotated, tracked)

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+supervision
+ultralytics
+timm
+scikit-image