diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..1cbcc7c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Joost Helfers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index 659945d..e3229ce 100644 --- a/README.md +++ b/README.md @@ -1 +1,104 @@ -# ComfyUI-CVOverlay \ No newline at end of file +# ComfyUI-OpenCV-Overlays + +This project integrates OpenCV and YOLOv8 into ComfyUI, providing custom nodes for object detection, blob tracking, and aesthetic video overlays. + +## Overview + +The `ComfyUI-CVOverlay` package provides four main custom nodes: + +- **CV Model Loader**: Loads YOLO models for object detection +- **CV Object Detector**: Performs real-time object detection using YOLOv8 +- **CV Blob Tracker**: Implements blob tracking algorithms using OpenCV +- **CV Aesthetic Overlay**: Applies customizable technical/surveillance-style overlays + +## Features + +✨ **Lightweight Integration**: Minimal dependencies, optimized for ComfyUI workflows +🎯 **Multiple Detection Models**: Support for YOLOv8n through YOLOv8x models +🎨 **Aesthetic Overlays**: Technical, surveillance, minimal, and cyberpunk styles +📹 **Video Processing**: Frame-by-frame processing for video workflows +🎛️ **Full Control**: Confidence thresholds, colors, opacity, line thickness + +## Installation + +### Method 1: Using ComfyUI-Manager (Recommended) +1. Install [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) if you haven't already +2. Open ComfyUI-Manager in ComfyUI +3. Go to "Install Custom Nodes" +4. Search for "ComfyUI-CVOverlay" and click Install +5. Restart ComfyUI + +### Method 2: Manual Installation +1. Navigate to your ComfyUI `custom_nodes` directory +2. Clone this repository: + ```bash + git clone https://github.com/joosthel/ComfyUI-CVOverlay.git + ``` +3. Dependencies will be automatically installed by ComfyUI-Manager on next startup +4. Restart ComfyUI + +**Note**: ComfyUI-Manager will automatically handle the installation of required dependencies (`opencv-python`, `ultralytics`, `torch`, etc.) when you restart ComfyUI. + +## Quick Start + +1. **Load Model**: Use `CV Model Loader` with `yolov8n.pt` (lightweight) +2. **Detect Objects**: Connect your image to `CV Object Detector` +3. **Apply Style**: Use `CV Aesthetic Overlay` to visualize detections +4. **Save Result**: Connect to ComfyUI's Save Image node + +See [USAGE.md](USAGE.md) for detailed examples and workflows. + +## Node Reference + +### CV Model Loader +- Loads YOLO models (n/s/m/l/x variants) +- Supports custom trained models +- Output: CV_MODEL + +### CV Object Detector +- Input: CV_MODEL, IMAGE +- Configurable confidence and IOU thresholds +- Output: IMAGE, CV_DETECTIONS + +### CV Blob Tracker +- Input: IMAGE, (optional) CV_TRACKS +- Multiple tracking algorithms +- Output: IMAGE, CV_TRACKS + +### CV Aesthetic Overlay +- Input: IMAGE, (optional) CV_DETECTIONS, CV_TRACKS +- Multiple overlay styles and color schemes +- Output: IMAGE + +## Project Structure + +``` +ComfyUI-CVOverlay/ +├── __init__.py # ComfyUI node registration +├── requirements.txt # Dependencies +├── USAGE.md # Usage examples and workflows +├── nodes/ # Custom nodes +│ ├── cv_model_loader.py +│ ├── cv_object_detector.py +│ ├── cv_blob_tracker.py +│ └── cv_aesthetic_overlay.py +└── utils/ # Helper functions + ├── opencv_helpers.py + └── yolo_utils.py +``` + +## Requirements + +- Python ≥ 3.8 +- ComfyUI +- OpenCV ≥ 4.8.0 +- Ultralytics (YOLOv8) ≥ 8.0.0 +- PyTorch ≥ 1.11.0 + +## Contributing + +Contributions welcome! Please submit pull requests or open issues for enhancements and bug fixes. + +## License + +MIT License - see LICENSE file for details. \ No newline at end of file diff --git a/USAGE.md b/USAGE.md new file mode 100644 index 0000000..e71dfdb --- /dev/null +++ b/USAGE.md @@ -0,0 +1,64 @@ +# Usage Examples + +## Installation & Setup + +1. Install via ComfyUI-Manager (search for "ComfyUI-CVOverlay") +2. Restart ComfyUI +3. Dependencies are installed automatically +4. Find nodes under "CV/" category in the node menu + +## Basic Object Detection Workflow + +1. **CV Model Loader** → Load a YOLO model (start with `yolov8n.pt` for lightweight processing) +2. **Load Image** (ComfyUI default node) → Your input image +3. **CV Object Detector** → Connect model and image, adjust confidence threshold +4. **CV Aesthetic Overlay** → Connect image and detections, choose style and colors +5. **Save Image** (ComfyUI default node) → Output the result + +## Basic Blob Tracking Workflow + +1. **Load Image** → Your input image/video frame +2. **CV Blob Tracker** → Process for moving objects +3. **CV Aesthetic Overlay** → Visualize tracked objects +4. **Save Image** → Output the result + +## Node Parameters Quick Reference + +### CV Model Loader +- `model_name`: Choose from yolov8n.pt (fastest) to yolov8x.pt (most accurate) +- `custom_model_path`: Use your own trained model (optional) + +### CV Object Detector +- `confidence`: 0.5 = balanced, 0.3 = more detections, 0.7 = fewer but more confident +- `iou_threshold`: 0.45 = default overlap filtering + +### CV Aesthetic Overlay +- `overlay_style`: + - `minimal` = Just corner markers + - `technical` = Full boxes with corners + - `surveillance` = Thick boxes with crosshairs + - `cyberpunk` = Glitched fragmented boxes +- `color_scheme`: green (classic), blue, red, white, orange, purple +- `opacity`: 0.8 = semi-transparent, 1.0 = fully opaque + +### CV Blob Tracker +- `tracking_method`: + - `background_subtraction` = Good for stationary camera + - `optical_flow` = Good for moving camera + - `contour_tracking` = Simple shape-based tracking +- `min_area`: Minimum size of objects to track (500 = medium sized objects) + +## Tips + +1. **Start Small**: Use `yolov8n.pt` model first, it's fast and good for testing +2. **Adjust Confidence**: Lower confidence (0.3-0.4) for more detections, higher (0.6-0.8) for precision +3. **Style Mixing**: Try different overlay styles with different color schemes for unique looks +4. **Video Processing**: For video, connect multiple frames through the same workflow +5. **Performance**: Disable unnecessary overlays (grid, crosshairs) for faster processing + +## Troubleshooting + +- **No detections**: Lower confidence threshold or try different YOLO model +- **Too many false positives**: Raise confidence threshold or IOU threshold +- **Tracking not working**: Try different tracking method or adjust min_area +- **Overlay too bright**: Reduce opacity or try different color scheme diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ab29714 --- /dev/null +++ b/__init__.py @@ -0,0 +1,67 @@ +""" +ComfyUI OpenCV Overlays - Custom nodes for computer vision effects +""" + +import os +import sys +import importlib.util + +# Add the current directory to the Python path +current_dir = os.path.dirname(os.path.abspath(__file__)) +nodes_dir = os.path.join(current_dir, "nodes") + +if current_dir not in sys.path: + sys.path.insert(0, current_dir) +if nodes_dir not in sys.path: + sys.path.insert(0, nodes_dir) + +# Import nodes with error handling using direct file imports +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +def import_node_from_file(file_path, class_name): + """Import a class from a specific file""" + try: + spec = importlib.util.spec_from_file_location("temp_module", file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return getattr(module, class_name) + except Exception as e: + print(f"Failed to import {class_name} from {file_path}: {e}") + return None + +# Import CV_ModelLoader +cv_model_loader_path = os.path.join(nodes_dir, "cv_model_loader.py") +if os.path.exists(cv_model_loader_path): + CV_ModelLoader = import_node_from_file(cv_model_loader_path, "CV_ModelLoader") + if CV_ModelLoader: + NODE_CLASS_MAPPINGS["CV_ModelLoader"] = CV_ModelLoader + NODE_DISPLAY_NAME_MAPPINGS["CV_ModelLoader"] = "CV Model Loader" + +# Import CV_ObjectDetector +cv_object_detector_path = os.path.join(nodes_dir, "cv_object_detector.py") +if os.path.exists(cv_object_detector_path): + CV_ObjectDetector = import_node_from_file(cv_object_detector_path, "CV_ObjectDetector") + if CV_ObjectDetector: + NODE_CLASS_MAPPINGS["CV_ObjectDetector"] = CV_ObjectDetector + NODE_DISPLAY_NAME_MAPPINGS["CV_ObjectDetector"] = "CV Object Detector" + +# Import CV_BlobTracker +cv_blob_tracker_path = os.path.join(nodes_dir, "cv_blob_tracker.py") +if os.path.exists(cv_blob_tracker_path): + CV_BlobTracker = import_node_from_file(cv_blob_tracker_path, "CV_BlobTracker") + if CV_BlobTracker: + NODE_CLASS_MAPPINGS["CV_BlobTracker"] = CV_BlobTracker + NODE_DISPLAY_NAME_MAPPINGS["CV_BlobTracker"] = "CV Blob Tracker" + +# Import CV_AestheticOverlay +cv_aesthetic_overlay_path = os.path.join(nodes_dir, "cv_aesthetic_overlay.py") +if os.path.exists(cv_aesthetic_overlay_path): + CV_AestheticOverlay = import_node_from_file(cv_aesthetic_overlay_path, "CV_AestheticOverlay") + if CV_AestheticOverlay: + NODE_CLASS_MAPPINGS["CV_AestheticOverlay"] = CV_AestheticOverlay + NODE_DISPLAY_NAME_MAPPINGS["CV_AestheticOverlay"] = "CV Aesthetic Overlay" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +print(f"ComfyUI-CVOverlay loaded {len(NODE_CLASS_MAPPINGS)} nodes successfully") \ No newline at end of file diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..74424af Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/install.bat b/install.bat new file mode 100644 index 0000000..256b3f5 --- /dev/null +++ b/install.bat @@ -0,0 +1,34 @@ +@echo off +echo Installing ComfyUI OpenCV Overlays dependencies... +echo. + +REM Navigate to the custom node directory +cd /d "%~dp0" + +REM Find ComfyUI's Python executable +set "COMFYUI_PYTHON=..\..\..\python_embeded\python.exe" +if not exist "%COMFYUI_PYTHON%" ( + set "COMFYUI_PYTHON=..\..\..\python\python.exe" +) +if not exist "%COMFYUI_PYTHON%" ( + set "COMFYUI_PYTHON=python" + echo Using system Python... +) else ( + echo Using ComfyUI's Python: %COMFYUI_PYTHON% +) + +echo. +echo Installing requirements... +"%COMFYUI_PYTHON%" -m pip install --upgrade pip +"%COMFYUI_PYTHON%" -m pip install -r requirements.txt + +echo. +echo Installing package in development mode... +"%COMFYUI_PYTHON%" -m pip install -e . + +echo. +echo Installation complete! +echo. +echo You can now restart ComfyUI to see the new CV nodes in the node menu. +echo The nodes will appear under the "CV/" category. +pause diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..f27a938 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1 @@ +# Nodes package diff --git a/nodes/__pycache__/__init__.cpython-312.pyc b/nodes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..2abf7e5 Binary files /dev/null and b/nodes/__pycache__/__init__.cpython-312.pyc differ diff --git a/nodes/__pycache__/cv_aesthetic_overlay.cpython-312.pyc b/nodes/__pycache__/cv_aesthetic_overlay.cpython-312.pyc new file mode 100644 index 0000000..9b4f2d0 Binary files /dev/null and b/nodes/__pycache__/cv_aesthetic_overlay.cpython-312.pyc differ diff --git a/nodes/__pycache__/cv_blob_tracker.cpython-312.pyc b/nodes/__pycache__/cv_blob_tracker.cpython-312.pyc new file mode 100644 index 0000000..c1d9e5e Binary files /dev/null and b/nodes/__pycache__/cv_blob_tracker.cpython-312.pyc differ diff --git a/nodes/__pycache__/cv_model_loader.cpython-312.pyc b/nodes/__pycache__/cv_model_loader.cpython-312.pyc new file mode 100644 index 0000000..d95b4f0 Binary files /dev/null and b/nodes/__pycache__/cv_model_loader.cpython-312.pyc differ diff --git a/nodes/__pycache__/cv_object_detector.cpython-312.pyc b/nodes/__pycache__/cv_object_detector.cpython-312.pyc new file mode 100644 index 0000000..2eb55b6 Binary files /dev/null and b/nodes/__pycache__/cv_object_detector.cpython-312.pyc differ diff --git a/nodes/cv_aesthetic_overlay.py b/nodes/cv_aesthetic_overlay.py new file mode 100644 index 0000000..92d8ba8 --- /dev/null +++ b/nodes/cv_aesthetic_overlay.py @@ -0,0 +1,261 @@ +""" +CV Aesthetic Overlay Node for ComfyUI +Applies technical/surveillance-style visual overlays to images with detection data +""" + +# Handle missing dependencies gracefully +try: + import cv2 + import torch + import numpy as np + from PIL import Image, ImageDraw, ImageFont + DEPENDENCIES_AVAILABLE = True +except ImportError as e: + DEPENDENCIES_AVAILABLE = False + MISSING_DEPS = str(e) + + +class CV_AestheticOverlay: + """Applies aesthetic overlays to images with detection/tracking data""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "overlay_style": (["minimal", "technical", "surveillance", "cyberpunk"], { + "default": "technical" + }), + "line_thickness": ("INT", { + "default": 2, + "min": 1, + "max": 10, + "step": 1 + }), + "opacity": ("FLOAT", { + "default": 0.8, + "min": 0.0, + "max": 1.0, + "step": 0.05 + }), + "color_scheme": (["green", "blue", "red", "white", "orange", "purple"], { + "default": "green" + }), + "show_coordinates": ("BOOLEAN", {"default": True}), + "show_confidence": ("BOOLEAN", {"default": True}), + "show_grid": ("BOOLEAN", {"default": False}), + "show_crosshairs": ("BOOLEAN", {"default": True}), + }, + "optional": { + "detections": ("CV_DETECTIONS",), + "tracks": ("CV_TRACKS",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_overlay" + CATEGORY = "CV/Overlay" + + def apply_overlay(self, image, overlay_style, line_thickness, opacity, color_scheme, + show_coordinates, show_confidence, show_grid, show_crosshairs, + detections=None, tracks=None): + """Apply aesthetic overlay to image""" + if not DEPENDENCIES_AVAILABLE: + raise RuntimeError(f"Missing dependencies: {MISSING_DEPS}. Please install requirements: pip install opencv-python torch Pillow") + + try: + # Convert ComfyUI image tensor to PIL Image + if isinstance(image, torch.Tensor): + if image.dim() == 4: + image = image[0] # Take first image from batch + img_np = (image.cpu().numpy() * 255).astype(np.uint8) + img_pil = Image.fromarray(img_np) + else: + img_pil = image + + # Create overlay + overlay = Image.new('RGBA', img_pil.size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay) + + # Color schemes + colors = { + 'green': (0, 255, 0), + 'blue': (0, 150, 255), + 'red': (255, 50, 50), + 'white': (255, 255, 255), + 'orange': (255, 165, 0), + 'purple': (180, 0, 255) + } + + primary_color = colors[color_scheme] + alpha = int(opacity * 255) + color_with_alpha = (*primary_color, alpha) + + # Draw grid if enabled + if show_grid: + self._draw_grid(draw, img_pil.size, color_with_alpha, line_thickness) + + # Draw detections + if detections: + for detection in detections: + self._draw_detection(draw, detection, overlay_style, color_with_alpha, + line_thickness, show_coordinates, show_confidence) + + # Draw tracks + if tracks: + for track in tracks: + self._draw_track(draw, track, overlay_style, color_with_alpha, + line_thickness, show_coordinates, show_confidence) + + # Draw crosshairs if enabled + if show_crosshairs: + self._draw_crosshairs(draw, img_pil.size, color_with_alpha, line_thickness) + + # Composite overlay onto original image + if img_pil.mode != 'RGBA': + img_pil = img_pil.convert('RGBA') + + result = Image.alpha_composite(img_pil, overlay) + result = result.convert('RGB') + + # Convert back to ComfyUI tensor format + result_np = np.array(result).astype(np.float32) / 255.0 + result_tensor = torch.from_numpy(result_np).unsqueeze(0) # Add batch dimension + + return (result_tensor,) + + except Exception as e: + raise RuntimeError(f"Overlay application failed: {str(e)}") + + def _draw_grid(self, draw, size, color, thickness): + """Draw coordinate grid""" + width, height = size + grid_spacing = min(width, height) // 20 + + # Vertical lines + for x in range(0, width, grid_spacing): + draw.line([(x, 0), (x, height)], fill=color, width=max(1, thickness//2)) + + # Horizontal lines + for y in range(0, height, grid_spacing): + draw.line([(0, y), (width, y)], fill=color, width=max(1, thickness//2)) + + def _draw_crosshairs(self, draw, size, color, thickness): + """Draw center crosshairs""" + width, height = size + center_x, center_y = width // 2, height // 2 + crosshair_size = min(width, height) // 20 + + # Horizontal crosshair + draw.line([(center_x - crosshair_size, center_y), + (center_x + crosshair_size, center_y)], fill=color, width=thickness) + + # Vertical crosshair + draw.line([(center_x, center_y - crosshair_size), + (center_x, center_y + crosshair_size)], fill=color, width=thickness) + + def _draw_detection(self, draw, detection, style, color, thickness, show_coords, show_conf): + """Draw detection overlay""" + bbox = detection['bbox'] + x1, y1, x2, y2 = bbox + + if style == "minimal": + # Just corner markers + corner_size = 10 + self._draw_corners(draw, bbox, color, thickness, corner_size) + + elif style == "technical": + # Rectangle + corner markers + info + draw.rectangle([x1, y1, x2, y2], outline=color, width=thickness) + self._draw_corners(draw, bbox, color, thickness, 15) + + elif style == "surveillance": + # Thick rectangle + crosshair at center + draw.rectangle([x1, y1, x2, y2], outline=color, width=thickness + 1) + center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2 + cross_size = 8 + draw.line([(center_x - cross_size, center_y), + (center_x + cross_size, center_y)], fill=color, width=thickness) + draw.line([(center_x, center_y - cross_size), + (center_x, center_y + cross_size)], fill=color, width=thickness) + + elif style == "cyberpunk": + # Glitch-style fragmented rectangle + self._draw_glitch_box(draw, bbox, color, thickness) + + # Add text information + if show_coords or show_conf: + text_y = max(y1 - 20, 10) + text_parts = [] + + if show_coords: + text_parts.append(f"({x1},{y1})") + + if show_conf and 'confidence' in detection: + conf = detection['confidence'] + text_parts.append(f"{conf:.2f}") + + if 'class_name' in detection: + text_parts.append(detection['class_name']) + + if text_parts: + text = " | ".join(text_parts) + draw.text((x1, text_y), text, fill=color) + + def _draw_track(self, draw, track, style, color, thickness, show_coords, show_conf): + """Draw tracking overlay""" + if 'bbox' in track: + # Draw like detection but with different style + fake_detection = { + 'bbox': track['bbox'], + 'confidence': track.get('confidence', 0.0) + } + self._draw_detection(draw, fake_detection, style, color, thickness, show_coords, show_conf) + + if 'center' in track: + # Draw center point + center_x, center_y = track['center'] + point_size = thickness + 2 + draw.ellipse([center_x - point_size, center_y - point_size, + center_x + point_size, center_y + point_size], + fill=color, outline=color) + + def _draw_corners(self, draw, bbox, color, thickness, corner_size): + """Draw corner markers on bounding box""" + x1, y1, x2, y2 = bbox + + # Top-left corner + draw.line([(x1, y1), (x1 + corner_size, y1)], fill=color, width=thickness) + draw.line([(x1, y1), (x1, y1 + corner_size)], fill=color, width=thickness) + + # Top-right corner + draw.line([(x2 - corner_size, y1), (x2, y1)], fill=color, width=thickness) + draw.line([(x2, y1), (x2, y1 + corner_size)], fill=color, width=thickness) + + # Bottom-left corner + draw.line([(x1, y2 - corner_size), (x1, y2)], fill=color, width=thickness) + draw.line([(x1, y2), (x1 + corner_size, y2)], fill=color, width=thickness) + + # Bottom-right corner + draw.line([(x2 - corner_size, y2), (x2, y2)], fill=color, width=thickness) + draw.line([(x2, y2 - corner_size), (x2, y2)], fill=color, width=thickness) + + def _draw_glitch_box(self, draw, bbox, color, thickness): + """Draw cyberpunk-style glitched bounding box""" + x1, y1, x2, y2 = bbox + + # Main rectangle with gaps + segments = [ + [(x1, y1), (x1 + (x2-x1)//3, y1)], # Top left segment + [(x1 + 2*(x2-x1)//3, y1), (x2, y1)], # Top right segment + [(x2, y1), (x2, y1 + (y2-y1)//3)], # Right top segment + [(x2, y1 + 2*(y2-y1)//3), (x2, y2)], # Right bottom segment + [(x2, y2), (x1 + 2*(x2-x1)//3, y2)], # Bottom right segment + [(x1 + (x2-x1)//3, y2), (x1, y2)], # Bottom left segment + [(x1, y2), (x1, y1 + 2*(y2-y1)//3)], # Left bottom segment + [(x1, y1 + (y2-y1)//3), (x1, y1)], # Left top segment + ] + + for segment in segments: + draw.line(segment, fill=color, width=thickness) diff --git a/nodes/cv_blob_tracker.py b/nodes/cv_blob_tracker.py new file mode 100644 index 0000000..cbd8a76 --- /dev/null +++ b/nodes/cv_blob_tracker.py @@ -0,0 +1,187 @@ +""" +CV Blob Tracker Node for ComfyUI +Tracks moving objects/blobs across video frames using OpenCV +""" + +# Handle missing dependencies gracefully +try: + import cv2 + import torch + import numpy as np + from PIL import Image + DEPENDENCIES_AVAILABLE = True +except ImportError as e: + DEPENDENCIES_AVAILABLE = False + MISSING_DEPS = str(e) + + +class CV_BlobTracker: + """Tracks blobs/objects across frames using OpenCV algorithms""" + + def __init__(self): + self.prev_frame = None + self.tracker_type = None + self.trackers = [] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "tracking_method": (["background_subtraction", "optical_flow", "contour_tracking"], { + "default": "background_subtraction" + }), + "min_area": ("INT", { + "default": 500, + "min": 50, + "max": 10000, + "step": 50 + }), + "threshold": ("INT", { + "default": 50, + "min": 10, + "max": 255, + "step": 5 + }), + }, + "optional": { + "previous_tracks": ("CV_TRACKS",), + } + } + + RETURN_TYPES = ("IMAGE", "CV_TRACKS") + RETURN_NAMES = ("image", "tracks") + FUNCTION = "track_blobs" + CATEGORY = "CV/Tracking" + + def track_blobs(self, image, tracking_method, min_area, threshold, previous_tracks=None): + """Track blobs in the current frame""" + if not DEPENDENCIES_AVAILABLE: + raise RuntimeError(f"Missing dependencies: {MISSING_DEPS}. Please install requirements: pip install opencv-python torch") + + try: + # Convert ComfyUI image tensor to OpenCV format + if isinstance(image, torch.Tensor): + if image.dim() == 4: + image = image[0] # Take first image from batch + img_np = (image.cpu().numpy() * 255).astype(np.uint8) + else: + img_np = np.array(image) + + # Convert to grayscale for processing + if len(img_np.shape) == 3: + gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) + else: + gray = img_np + + tracks = [] + + if tracking_method == "background_subtraction": + tracks = self._background_subtraction_tracking(gray, min_area, threshold) + elif tracking_method == "optical_flow": + tracks = self._optical_flow_tracking(gray, previous_tracks) + elif tracking_method == "contour_tracking": + tracks = self._contour_tracking(gray, min_area, threshold) + + # Store current frame for next iteration + self.prev_frame = gray.copy() + + # Convert back to ComfyUI tensor format + if isinstance(image, torch.Tensor): + output_image = image + else: + img_tensor = torch.from_numpy(img_np).float() / 255.0 + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) + output_image = img_tensor + + return (output_image, tracks) + + except Exception as e: + raise RuntimeError(f"Blob tracking failed: {str(e)}") + + def _background_subtraction_tracking(self, gray, min_area, threshold): + """Simple background subtraction tracking""" + tracks = [] + + if self.prev_frame is not None: + # Calculate frame difference + diff = cv2.absdiff(self.prev_frame, gray) + _, thresh = cv2.threshold(diff, threshold, 255, cv2.THRESH_BINARY) + + # Find contours + contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + for contour in contours: + area = cv2.contourArea(contour) + if area > min_area: + x, y, w, h = cv2.boundingRect(contour) + center_x = x + w // 2 + center_y = y + h // 2 + + tracks.append({ + 'bbox': [x, y, x + w, y + h], + 'center': [center_x, center_y], + 'area': area, + 'confidence': min(area / (min_area * 10), 1.0) + }) + + return tracks + + def _optical_flow_tracking(self, gray, previous_tracks): + """Lucas-Kanade optical flow tracking""" + tracks = [] + + if self.prev_frame is not None and previous_tracks: + # Extract previous points + prev_points = [] + for track in previous_tracks: + if 'center' in track: + prev_points.append([track['center']]) + + if prev_points: + prev_points = np.array(prev_points, dtype=np.float32) + + # Calculate optical flow + new_points, status, error = cv2.calcOpticalFlowPyrLK( + self.prev_frame, gray, prev_points, None + ) + + # Keep only good points + good_new = new_points[status == 1] + + for i, point in enumerate(good_new): + x, y = point.ravel() + tracks.append({ + 'center': [int(x), int(y)], + 'confidence': 0.8, + 'bbox': [int(x-20), int(y-20), int(x+20), int(y+20)] # Approximate bbox + }) + + return tracks + + def _contour_tracking(self, gray, min_area, threshold): + """Simple contour-based tracking""" + tracks = [] + + # Apply threshold + _, thresh = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) + + # Find contours + contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + for contour in contours: + area = cv2.contourArea(contour) + if area > min_area: + x, y, w, h = cv2.boundingRect(contour) + center_x = x + w // 2 + center_y = y + h // 2 + + tracks.append({ + 'bbox': [x, y, x + w, y + h], + 'center': [center_x, center_y], + 'area': area, + 'confidence': min(area / (min_area * 5), 1.0) + }) + + return tracks diff --git a/nodes/cv_model_loader.py b/nodes/cv_model_loader.py new file mode 100644 index 0000000..887a6eb --- /dev/null +++ b/nodes/cv_model_loader.py @@ -0,0 +1,57 @@ +""" +CV Model Loader Node for ComfyUI +Loads YOLO models for object detection +""" + +import os + +# Handle missing dependencies gracefully +try: + import torch + from ultralytics import YOLO + DEPENDENCIES_AVAILABLE = True +except ImportError as e: + DEPENDENCIES_AVAILABLE = False + MISSING_DEPS = str(e) + + +class CV_ModelLoader: + """Loads YOLO models for computer vision tasks""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_name": (["yolov8n.pt", "yolov8s.pt", "yolov8m.pt", "yolov8l.pt", "yolov8x.pt"], { + "default": "yolov8n.pt" + }), + }, + "optional": { + "custom_model_path": ("STRING", { + "default": "", + "multiline": False + }), + } + } + + RETURN_TYPES = ("CV_MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "load_model" + CATEGORY = "CV/Models" + + def load_model(self, model_name, custom_model_path=""): + """Load YOLO model""" + if not DEPENDENCIES_AVAILABLE: + raise RuntimeError(f"Missing dependencies: {MISSING_DEPS}. Please install requirements: pip install torch ultralytics") + + try: + # Use custom path if provided, otherwise use default model + model_path = custom_model_path if custom_model_path else model_name + + # Load YOLO model + model = YOLO(model_path) + + return (model,) + + except Exception as e: + raise RuntimeError(f"Failed to load model: {str(e)}") diff --git a/nodes/cv_object_detector.py b/nodes/cv_object_detector.py new file mode 100644 index 0000000..8495b9c --- /dev/null +++ b/nodes/cv_object_detector.py @@ -0,0 +1,95 @@ +""" +CV Object Detector Node for ComfyUI +Performs object detection using YOLO models +""" + +# Handle missing dependencies gracefully +try: + import torch + import numpy as np + from PIL import Image + DEPENDENCIES_AVAILABLE = True +except ImportError as e: + DEPENDENCIES_AVAILABLE = False + MISSING_DEPS = str(e) + + +class CV_ObjectDetector: + """Detects objects in images using YOLO models""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("CV_MODEL",), + "image": ("IMAGE",), + "confidence": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "iou_threshold": ("FLOAT", { + "default": 0.45, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + } + } + + RETURN_TYPES = ("IMAGE", "CV_DETECTIONS") + RETURN_NAMES = ("image", "detections") + FUNCTION = "detect_objects" + CATEGORY = "CV/Detection" + + def detect_objects(self, model, image, confidence, iou_threshold): + """Perform object detection on input image""" + if not DEPENDENCIES_AVAILABLE: + raise RuntimeError(f"Missing dependencies: {MISSING_DEPS}. Please install requirements: pip install torch ultralytics opencv-python") + + try: + # Convert ComfyUI image tensor to PIL Image + if isinstance(image, torch.Tensor): + # ComfyUI images are in format (batch, height, width, channels) + if image.dim() == 4: + image = image[0] # Take first image from batch + # Convert from tensor to numpy + img_np = (image.cpu().numpy() * 255).astype(np.uint8) + img_pil = Image.fromarray(img_np) + else: + img_pil = image + + # Run inference + results = model(img_pil, conf=confidence, iou=iou_threshold) + + # Extract detection data + detections = [] + if len(results) > 0 and results[0].boxes is not None: + boxes = results[0].boxes + for i in range(len(boxes)): + box = boxes.xyxy[i].cpu().numpy() # [x1, y1, x2, y2] + conf = boxes.conf[i].cpu().numpy() + cls = int(boxes.cls[i].cpu().numpy()) + + detections.append({ + 'bbox': box.tolist(), + 'confidence': float(conf), + 'class': cls, + 'class_name': model.names[cls] if hasattr(model, 'names') else str(cls) + }) + + # Convert back to ComfyUI tensor format + if isinstance(image, torch.Tensor): + output_image = image + else: + # Convert PIL to tensor if needed + img_tensor = torch.from_numpy(np.array(img_pil)).float() / 255.0 + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + output_image = img_tensor + + return (output_image, detections) + + except Exception as e: + raise RuntimeError(f"Object detection failed: {str(e)}") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..40e0f62 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +opencv-python +ultralytics +torch +torchvision +numpy +Pillow \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..dd7ee44 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1 @@ +# Utils package diff --git a/utils/opencv_helpers.py b/utils/opencv_helpers.py new file mode 100644 index 0000000..e49ae6a --- /dev/null +++ b/utils/opencv_helpers.py @@ -0,0 +1,103 @@ +""" +Utility functions for OpenCV operations +""" + +import cv2 +import numpy as np +import torch +from PIL import Image + + +def tensor_to_opencv(tensor): + """Convert ComfyUI tensor to OpenCV format""" + if isinstance(tensor, torch.Tensor): + if tensor.dim() == 4: + tensor = tensor[0] # Take first image from batch + img_np = (tensor.cpu().numpy() * 255).astype(np.uint8) + # Convert RGB to BGR for OpenCV + if len(img_np.shape) == 3: + img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + return img_np + return tensor + + +def opencv_to_tensor(img_np): + """Convert OpenCV format to ComfyUI tensor""" + # Convert BGR to RGB + if len(img_np.shape) == 3: + img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB) + + # Normalize and convert to tensor + img_tensor = torch.from_numpy(img_np).float() / 255.0 + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + + return img_tensor + + +def tensor_to_pil(tensor): + """Convert ComfyUI tensor to PIL Image""" + if isinstance(tensor, torch.Tensor): + if tensor.dim() == 4: + tensor = tensor[0] # Take first image from batch + img_np = (tensor.cpu().numpy() * 255).astype(np.uint8) + return Image.fromarray(img_np) + return tensor + + +def pil_to_tensor(img_pil): + """Convert PIL Image to ComfyUI tensor""" + img_np = np.array(img_pil).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_np) + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + return img_tensor + + +def normalize_bbox(bbox, img_width, img_height): + """Normalize bounding box coordinates to 0-1 range""" + x1, y1, x2, y2 = bbox + return [ + x1 / img_width, + y1 / img_height, + x2 / img_width, + y2 / img_height + ] + + +def denormalize_bbox(bbox, img_width, img_height): + """Denormalize bounding box coordinates from 0-1 range to pixel coordinates""" + x1, y1, x2, y2 = bbox + return [ + int(x1 * img_width), + int(y1 * img_height), + int(x2 * img_width), + int(y2 * img_height) + ] + + +def resize_image_keep_aspect(image, target_size): + """Resize image while keeping aspect ratio""" + h, w = image.shape[:2] + target_w, target_h = target_size + + # Calculate scale factor + scale = min(target_w / w, target_h / h) + + # Calculate new dimensions + new_w = int(w * scale) + new_h = int(h * scale) + + # Resize image + resized = cv2.resize(image, (new_w, new_h)) + + # Create padded image + result = np.zeros((target_h, target_w, 3), dtype=image.dtype) + + # Center the resized image + y_offset = (target_h - new_h) // 2 + x_offset = (target_w - new_w) // 2 + + result[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized + + return result, scale, (x_offset, y_offset) diff --git a/utils/yolo_utils.py b/utils/yolo_utils.py new file mode 100644 index 0000000..6f28e20 --- /dev/null +++ b/utils/yolo_utils.py @@ -0,0 +1,132 @@ +""" +Utility functions for YOLO operations +""" + +import torch +import numpy as np +from ultralytics import YOLO + + +def load_yolo_model(model_path): + """Load YOLO model safely""" + try: + model = YOLO(model_path) + return model + except Exception as e: + raise RuntimeError(f"Failed to load YOLO model from {model_path}: {str(e)}") + + +def filter_detections_by_class(detections, allowed_classes): + """Filter detections by class names or IDs""" + if not allowed_classes: + return detections + + filtered = [] + for detection in detections: + class_id = detection.get('class', -1) + class_name = detection.get('class_name', '') + + if class_id in allowed_classes or class_name in allowed_classes: + filtered.append(detection) + + return filtered + + +def filter_detections_by_confidence(detections, min_confidence): + """Filter detections by confidence threshold""" + return [d for d in detections if d.get('confidence', 0) >= min_confidence] + + +def filter_detections_by_area(detections, min_area=None, max_area=None): + """Filter detections by bounding box area""" + filtered = [] + + for detection in detections: + bbox = detection['bbox'] + x1, y1, x2, y2 = bbox + area = (x2 - x1) * (y2 - y1) + + if min_area is not None and area < min_area: + continue + if max_area is not None and area > max_area: + continue + + filtered.append(detection) + + return filtered + + +def nms_detections(detections, iou_threshold=0.45): + """Apply Non-Maximum Suppression to detections""" + if not detections: + return detections + + # Convert to format for NMS + boxes = [] + scores = [] + + for detection in detections: + bbox = detection['bbox'] + boxes.append(bbox) + scores.append(detection.get('confidence', 1.0)) + + boxes = torch.tensor(boxes, dtype=torch.float32) + scores = torch.tensor(scores, dtype=torch.float32) + + # Apply NMS + keep_indices = torch.ops.torchvision.nms(boxes, scores, iou_threshold) + + # Return filtered detections + return [detections[i] for i in keep_indices] + + +def get_yolo_class_names(): + """Get standard YOLO class names""" + return [ + 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', + 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', + 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', + 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', + 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', + 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', + 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', + 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', + 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', + 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', + 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', + 'toothbrush' + ] + + +def convert_yolo_to_detection_format(results, img_width, img_height): + """Convert YOLO results to standardized detection format""" + detections = [] + + if len(results) > 0 and results[0].boxes is not None: + boxes = results[0].boxes + + for i in range(len(boxes)): + # Get box coordinates + box = boxes.xyxy[i].cpu().numpy() # [x1, y1, x2, y2] + conf = float(boxes.conf[i].cpu().numpy()) + cls = int(boxes.cls[i].cpu().numpy()) + + # Get class name if available + class_names = get_yolo_class_names() + class_name = class_names[cls] if cls < len(class_names) else f"class_{cls}" + + # Create detection dictionary + detection = { + 'bbox': box.tolist(), + 'confidence': conf, + 'class': cls, + 'class_name': class_name, + 'center': [(box[0] + box[2]) / 2, (box[1] + box[3]) / 2], + 'width': box[2] - box[0], + 'height': box[3] - box[1], + 'area': (box[2] - box[0]) * (box[3] - box[1]) + } + + detections.append(detection) + + return detections