add - initial setup for object detection
This commit is contained in:
joosthel
@@ -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.
|
||||
@@ -1 +1,104 @@
|
||||
# ComfyUI-CVOverlay
|
||||
# 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.
|
||||
@@ -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
|
||||
+67
@@ -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")
|
||||
Binary file not shown.
+34
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
# Nodes package
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)}")
|
||||
@@ -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)}")
|
||||
@@ -0,0 +1,6 @@
|
||||
opencv-python
|
||||
ultralytics
|
||||
torch
|
||||
torchvision
|
||||
numpy
|
||||
Pillow
|
||||
@@ -0,0 +1 @@
|
||||
# Utils package
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user