diff --git a/.gitignore b/.gitignore index c1c6240..b825e20 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ build/ .eggs/ test_output/ output/ +models/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..3d9946b --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,242 @@ +# AGENTS.md — ComfyUI-VideoActorExtract + +Guidelines for agentic coding agents working in this repository. + +## Project Overview + +ComfyUI custom node plugin that detects, tracks, and extracts individual actors from video, +outputting per-actor green-screen videos and structured JSON metadata. + +- **Language**: Python 3.10+ (backend), JavaScript (ComfyUI widget) +- **Framework**: ComfyUI custom node plugin +- **Runtime**: Depends on ComfyUI host process; not a standalone application + +## Repository Structure + +``` +__init__.py # Plugin entry — dynamic importlib loading to avoid ComfyUI namespace conflicts +core/ + config.py # Global constants (thresholds, defaults, codec settings) + video_reader.py # Video info extraction and frame sampling +nodes/ + actor_extractor.py # Main ComfyUI node class (VideoActorExtractor) — pipeline orchestration +pipeline/ + detector.py # YOLOv8 person detection → BoundingBox list + segmenter.py # YOLOv8-seg person segmentation → mask per person + mask_tracker.py # Centroid-distance mask tracking across frames + tracker.py # ByteTrack multi-object tracking (optional, fallback to greedy) + identity.py # InsightFace face embedding clustering for same-actor merging + cropper.py # Green-screen compositing per detected person + merger.py # FFmpeg encoding + JSON metadata generation +js/ + widget.js # ComfyUI frontend extension (JSON display button) +``` + +## Build & Install + +```bash +# Install dependencies +pip install -r requirements.txt + +# Install in dev mode (optional) +pip install -e ".[dev]" +``` + +No build step required — this is a pure-Python ComfyUI plugin loaded at runtime. + +## Linting + +```bash +# Lint with ruff (uses default rules — no ruff.toml/config exists) +ruff check . + +# Auto-fix where possible +ruff check --fix . +``` + +No formatter is configured. Keep existing style when making changes. + +## Testing + +```bash +# Run all tests +pytest + +# Run a single test file +pytest tests/test_foo.py + +# Run a specific test by name +pytest tests/test_foo.py::test_specific_function -v +``` + +> **Note**: No test files exist yet. `pytest` is a dev dependency in `pyproject.toml`. +> Integration testing is done manually via ComfyUI's HTTP API (see README.md lines 283–367). +> When adding tests, create a `tests/` directory at the project root. + +## Code Style Guidelines + +### Naming Conventions (PEP 8) + +| Element | Style | Example | +|---------|-------|---------| +| Functions / methods | `snake_case` | `_build_continuous_segments`, `detect_masks` | +| Variables | `snake_case` | `actor_frames`, `mask_resized` | +| Classes | `PascalCase` | `VideoActorExtractor`, `MaskTracker` | +| Constants | `UPPER_SNAKE_CASE` | `DEFAULT_MAX_ACTORS`, `GREEN_SCREEN_COLOR` | +| Private members | `_leading_underscore` | `_ensure_loaded`, `_compute_centroid` | +| ComfyUI node attributes | `UPPER_SNAKE_CASE` | `INPUT_TYPES`, `RETURN_TYPES`, `FUNCTION` | + +### Import Order + +Group imports in this order, separated by blank lines: + +```python +# 1. Standard library +import os +import sys +import json + +# 2. Third-party +import torch +import numpy as np +from typing import List, Tuple, Optional + +# 3. Local project (absolute paths from project root) +from core.config import DEFAULT_MAX_ACTORS +from pipeline.detector import PersonDetector +``` + +**Lazy imports for heavy dependencies** — load YOLO/InsightFace inside methods on first use, +not at module top level. This keeps plugin import time fast: + +```python +def _ensure_loaded(self): + """Lazy-load InsightFace model.""" + if self._loaded: + return + import insightface + ... +``` + +### Type Annotations + +- Use `typing.List`, `typing.Dict`, `typing.Tuple`, `typing.Optional` (Python 3.10 compat). +- Annotate public function signatures; private helpers may omit annotations for brevity. +- Use `@dataclass` for data containers (e.g., `FrameRecord`, `VideoInfo`). + +```python +def detect(self, frame: np.ndarray) -> List[BoundingBox]: + """Detect all persons in a single frame.""" + ... +``` + +### Docstrings (Google style) + +```python +def _build_continuous_segments( + actor_frames: List[Tuple[int, np.ndarray, int]], + max_gap: int = 30, +) -> Tuple[List[List[np.ndarray]], List[dict]]: + """ + Group frames into continuous segments with gap interpolation. + + Args: + actor_frames: List of (frame_idx, frame_array, mask_id) tuples. + max_gap: Maximum gap in frames before starting a new segment. + + Returns: + Tuple of (segmented_frame_lists, segment_metadata). + """ +``` + +- Module-level docstrings: single-line summary. +- Class docstrings: one-line description of purpose. +- Public methods: `Args` / `Returns` sections. +- Skip docstrings for trivial private helpers. + +### Logging + +Use `print()` with a bracketed prefix — no `logging` module: + +```python +print("[VideoActorExtract] Step 1: Initializing person segmenter...") +print("[MaskTracker] Lost track for actor_id=3 at frame 142") +``` + +Module prefix convention: `[VideoActorExtract]`, `[PersonSegmenter]`, `[MaskTracker]`, `[Identity]`. + +### Error Handling + +- **Graceful degradation**: wrap optional heavy imports in `try/except ImportError`, + provide fallback behavior (e.g., greedy tracker when ByteTrack is unavailable). +- **No hard crashes in pipeline**: return empty/error JSON rather than raising from the + main `extract()` method. +- **ValueError** for invalid user inputs (bad video path, missing frames). +- Silent `except Exception: pass` is acceptable only for truly optional features (e.g., proxy setup). + +### Data Containers + +Prefer `@dataclass` for structured data; use `__slots__` for performance-critical containers: + +```python +@dataclass +class FrameRecord: + frame_idx: int + x1: float + y1: float + x2: float + y2: float + +class BoundingBox: + __slots__ = ['x1', 'y1', 'x2', 'y2', 'confidence'] +``` + +### ComfyUI Node Pattern + +All nodes must expose these class attributes: + +```python +class MyNode: + @classmethod + def INPUT_TYPES(cls): + return {"required": {...}, "optional": {...}} + + RETURN_TYPES = ("STRING", "STRING", "IMAGE") + RETURN_NAMES = ("result_json", "output_dir", "preview") + FUNCTION = "process" # method name for ComfyUI to call + CATEGORY = "video/actor" + OUTPUT_NODE = True # optional +``` + +### JavaScript (ComfyUI Widget) + +Minimal — single file in `js/`. Register via: + +```javascript +app.registerExtension({ + name: "VideoActorExtract.widget", + // ... +}); +``` + +## Commit Messages + +Conventional commits, lowercase prefix: + +``` +feat: add person segmentation pipeline +fix: resolve namespace conflict with ComfyUI nodes module +docs: comprehensive README with architecture guide +improve: identity clustering with connected components +chore: add .gitignore +``` + +## Key Constraints + +1. **No standalone execution**: This plugin only runs inside ComfyUI. There is no CLI entry point. +2. **`__init__.py` uses importlib**: All modules are loaded via `importlib.util` with prefixed names + (`vae_*`) to avoid namespace collisions with ComfyUI's own `nodes` module. Do not switch to + regular imports in `__init__.py`. +3. **Model weights** (`*.pt`) are gitignored. Models are loaded from `ComfyUI/models/video-actor-extract/`. +4. **Apple Silicon**: `onnxruntime-silicon` is used on ARM64 macOS (see requirements.txt conditional). +5. **Memory-sensitive**: Process frames one at a time; avoid accumulating full-video tensors in memory. diff --git a/__init__.py b/__init__.py index bca4726..b5c69be 100644 --- a/__init__.py +++ b/__init__.py @@ -15,6 +15,7 @@ if plugin_dir not in sys.path: # Use importlib to avoid namespace conflicts with ComfyUI's own 'nodes' module import importlib.util + def _load_module(name, path): spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) @@ -22,29 +23,52 @@ def _load_module(name, path): spec.loader.exec_module(module) return module + # Load core modules first -core_config = _load_module("vae_core_config", os.path.join(plugin_dir, "core", "config.py")) -vae_video_reader = _load_module("vae_video_reader", os.path.join(plugin_dir, "core", "video_reader.py")) +core_config = _load_module( + "vae_core_config", os.path.join(plugin_dir, "core", "config.py") +) +vae_video_reader = _load_module( + "vae_video_reader", os.path.join(plugin_dir, "core", "video_reader.py") +) # Load pipeline modules -vae_detector = _load_module("vae_detector", os.path.join(plugin_dir, "pipeline", "detector.py")) -vae_tracker = _load_module("vae_tracker", os.path.join(plugin_dir, "pipeline", "tracker.py")) -vae_identity = _load_module("vae_identity", os.path.join(plugin_dir, "pipeline", "identity.py")) -vae_cropper = _load_module("vae_cropper", os.path.join(plugin_dir, "pipeline", "cropper.py")) -vae_merger = _load_module("vae_merger", os.path.join(plugin_dir, "pipeline", "merger.py")) -vae_segmenter = _load_module("vae_segmenter", os.path.join(plugin_dir, "pipeline", "segmenter.py")) +vae_detector = _load_module( + "vae_detector", os.path.join(plugin_dir, "pipeline", "detector.py") +) +vae_tracker = _load_module( + "vae_tracker", os.path.join(plugin_dir, "pipeline", "tracker.py") +) +vae_identity = _load_module( + "vae_identity", os.path.join(plugin_dir, "pipeline", "identity.py") +) +vae_cropper = _load_module( + "vae_cropper", os.path.join(plugin_dir, "pipeline", "cropper.py") +) +vae_merger = _load_module( + "vae_merger", os.path.join(plugin_dir, "pipeline", "merger.py") +) +vae_segmenter = _load_module( + "vae_segmenter", os.path.join(plugin_dir, "pipeline", "segmenter.py") +) # Load main node -vae_actor_extractor = _load_module("vae_actor_extractor", os.path.join(plugin_dir, "nodes", "actor_extractor.py")) +vae_actor_extractor = _load_module( + "vae_actor_extractor", os.path.join(plugin_dir, "nodes", "actor_extractor.py") +) +vae_select_preview = _load_module( + "vae_select_preview", os.path.join(plugin_dir, "nodes", "select_preview.py") +) NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} -# Register nodes from the actor extractor module -if hasattr(vae_actor_extractor, "NODE_CLASS_MAPPINGS"): - NODE_CLASS_MAPPINGS.update(vae_actor_extractor.NODE_CLASS_MAPPINGS) -if hasattr(vae_actor_extractor, "NODE_DISPLAY_NAME_MAPPINGS"): - NODE_DISPLAY_NAME_MAPPINGS.update(vae_actor_extractor.NODE_DISPLAY_NAME_MAPPINGS) +# Register nodes from all node modules +for mod in [vae_actor_extractor, vae_select_preview]: + if hasattr(mod, "NODE_CLASS_MAPPINGS"): + NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS) + if hasattr(mod, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS) WEB_DIRECTORY = "./js" __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/nodes/actor_extractor.py b/nodes/actor_extractor.py index 8a5c065..67064c2 100644 --- a/nodes/actor_extractor.py +++ b/nodes/actor_extractor.py @@ -28,15 +28,13 @@ from typing import Tuple, List, Dict, Optional # Add parent directory to path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from core.video_reader import get_video_info, extract_frames +from core.video_reader import get_video_info from core.config import ( DEFAULT_MAX_ACTORS, DEFAULT_FACE_THRESHOLD, - DEFAULT_FPS_SAMPLE, DEFAULT_MIN_TRACK_LENGTH, DEFAULT_YOLO_MODEL, DEFAULT_MAX_LOST_FRAMES, - GREEN_SCREEN_COLOR, SEGMENT_GAP_SEC, ) @@ -45,13 +43,10 @@ try: from core.config import DEFAULT_SEG_MODEL except ImportError: DEFAULT_SEG_MODEL = "yolov8n-seg.pt" -from pipeline.detector import PersonDetector, BoundingBox -from pipeline.tracker import PersonTracker from pipeline.identity import IdentityCluster # face_threshold is kept for backward compatibility but identity clustering # now uses spatial-temporal constraints to prevent merging co-occurring tracks. -from pipeline.cropper import ActorCropper from pipeline.merger import merge_segments, generate_actor_json # PersonSegmenter may not exist yet (parallel task creating it) @@ -259,8 +254,7 @@ class VideoActorExtractor: Output: - actor_info_json: JSON string with actor information - output_dir: path to output directory containing per-actor MP4s - - actor_preview_images: IMAGE tensor of shape [num_actors, 5, H, W, 3], - top 5 preview frames per actor (padded with black frames if < 5 available) + - actor_count: number of detected actors """ @classmethod @@ -307,8 +301,8 @@ class VideoActorExtractor: }, } - RETURN_TYPES = ("STRING", "STRING", "IMAGE") - RETURN_NAMES = ("actor_info_json", "output_dir", "actor_preview_images") + RETURN_TYPES = ("STRING", "STRING", "INT") + RETURN_NAMES = ("actor_info_json", "output_dir", "actor_count") FUNCTION = "extract" CATEGORY = "video/actor" @@ -321,7 +315,7 @@ class VideoActorExtractor: max_actors: int = DEFAULT_MAX_ACTORS, face_threshold: float = DEFAULT_FACE_THRESHOLD, min_track_length: int = DEFAULT_MIN_TRACK_LENGTH, - ) -> Tuple[str, str, torch.Tensor]: + ) -> Tuple[str, str, int]: """ Run the full actor extraction pipeline. @@ -344,7 +338,7 @@ class VideoActorExtractor: min_track_length: Minimum number of frames for a valid track Returns: - (actor_info_json, output_dir, actor_preview_images) + (actor_info_json, output_dir, actor_count) """ output_dir = _get_comfyui_output_dir() @@ -458,8 +452,7 @@ class VideoActorExtractor: }, indent=2, ) - empty_preview = torch.zeros(0, 5, img_h, img_w, 3, dtype=torch.float32) - return (empty_json, output_dir, empty_preview) + return (empty_json, output_dir, 0) segmenter = PersonSegmenter(model_path=seg_model_path) @@ -519,8 +512,7 @@ class VideoActorExtractor: }, indent=2, ) - empty_preview = torch.zeros(0, 5, img_h, img_w, 3, dtype=torch.float32) - return (empty_json, output_dir, empty_preview) + return (empty_json, output_dir, 0) # ---------------------------------------------------------------- # Step 5: Filter short actors @@ -549,8 +541,7 @@ class VideoActorExtractor: }, indent=2, ) - empty_preview = torch.zeros(0, 5, img_h, img_w, 3, dtype=torch.float32) - return (empty_json, output_dir, empty_preview) + return (empty_json, output_dir, 0) # ---------------------------------------------------------------- # Step 5a: Split mixed tracks (tracks containing different people) @@ -620,7 +611,25 @@ class VideoActorExtractor: # be the same person. This prevents merging different people who # happen to have similar faces (e.g., twins, or face embedding limits). print("[VideoActorExtract] Step 5: Clustering actor identities...") - identity = IdentityCluster(threshold=face_threshold) + + # Resolve model directory for InsightFace (same location as YOLO models) + insightface_model_dir = "" + try: + import folder_paths + + insightface_model_dir = folder_paths.get_folder_paths("video-actor-extract") + insightface_model_dir = ( + insightface_model_dir[0] if insightface_model_dir else "" + ) + except Exception: + pass + if not insightface_model_dir: + # Fallback: derive from the resolved seg_model_path's parent directory + insightface_model_dir = os.path.dirname(seg_model_path) + + identity = IdentityCluster( + threshold=face_threshold, model_dir=insightface_model_dir + ) # Build synthetic FrameRecords from mask actors for identity clustering actor_synthetic_records: Dict[int, list] = {} @@ -658,9 +667,11 @@ class VideoActorExtractor: "duration_sec": total_frames / fps if fps > 0 else 0, } actor_data_for_json = {} - actor_preview_groups: List[List[np.ndarray]] = [] + preview_dir = os.path.join(output_dir, "previews") + os.makedirs(preview_dir, exist_ok=True) + import cv2 - for actor_id in actor_ids_sorted: + for i, actor_id in enumerate(actor_ids_sorted): # Merge all frames from all mask tracks belonging to this identity actor_all_frames: List[Tuple[int, np.ndarray, int]] = [] for aid in identity_to_actors[actor_id]: @@ -680,18 +691,11 @@ class VideoActorExtractor: key=lambda x: x[0], ) - # Select top 5 preview frames by mask area (descending) + # Save top 5 preview frames by mask area (descending) to disk top5 = sorted(actor_all_frames, key=lambda x: x[2], reverse=True)[:5] - actor_frames: List[np.ndarray] = [] - for _, masked_bgr, _ in top5: - # Convert BGR uint8 -> RGB float32 normalized 0-1 - rgb = masked_bgr[:, :, ::-1].copy() - rgb_float = rgb.astype(np.float32) / 255.0 - actor_frames.append(rgb_float) - # Pad to exactly 5 frames if fewer available - while len(actor_frames) < 5: - actor_frames.append(np.zeros((img_h, img_w, 3), dtype=np.float32)) - actor_preview_groups.append(actor_frames) + for k, (_, masked_bgr, _) in enumerate(top5): + preview_path = os.path.join(preview_dir, f"actor_{i}_{k}.jpg") + cv2.imwrite(preview_path, masked_bgr) # Build continuous segments with interpolation segments_frames, segments_info = _build_continuous_segments( @@ -724,21 +728,7 @@ class VideoActorExtractor: } # ---------------------------------------------------------------- - # Step 8: Build preview tensor - # ---------------------------------------------------------------- - if actor_preview_groups: - # Shape: [num_actors, 5, H, W, 3] - preview_tensor = torch.from_numpy( - np.stack( - [np.stack(group, axis=0) for group in actor_preview_groups], axis=0 - ) - ) - print(f"[VideoActorExtract] Preview tensor shape: {preview_tensor.shape}") - else: - preview_tensor = torch.zeros(0, 5, img_h, img_w, 3, dtype=torch.float32) - - # ---------------------------------------------------------------- - # Step 9: Generate JSON + # Step 8: Generate JSON # ---------------------------------------------------------------- json_path = os.path.join(output_dir, "actor_info.json") generate_actor_json(actor_data_for_json, video_info_dict, json_path) @@ -747,7 +737,7 @@ class VideoActorExtractor: json_content = f.read() print(f"[VideoActorExtract] Done! Output in {output_dir}") - return (json_content, output_dir, preview_tensor) + return (json_content, output_dir, len(actor_ids_sorted)) # ComfyUI node registration diff --git a/nodes/select_preview.py b/nodes/select_preview.py new file mode 100644 index 0000000..7b34777 --- /dev/null +++ b/nodes/select_preview.py @@ -0,0 +1,73 @@ +"""SelectActorPreview node for ComfyUI. + +Selects a specific preview image for a specific actor from the +preview files saved by VideoActorExtractor. +""" + +import os +import numpy as np +import torch +from typing import Tuple + + +class SelectActorPreview: + """Select a specific preview image for a specific actor.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "output_dir": ("STRING", {"default": ""}), + "actor_index": ("INT", {"default": 0, "min": 0, "max": 999}), + "photo_index": ("INT", {"default": 0, "min": 0, "max": 4}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "select" + CATEGORY = "video/actor" + OUTPUT_NODE = False + + def select( + self, output_dir: str, actor_index: int, photo_index: int + ) -> Tuple[torch.Tensor]: + """Select a preview image from disk for a specific actor. + + Args: + output_dir: Output directory from VideoActorExtractor. + actor_index: Which actor (0-based). + photo_index: Which photo (0-4). + + Returns: + Tuple of one 4D tensor [1, H, W, 3] RGB float32 in 0-1 range. + Must return a tuple (not bare tensor) so ComfyUI's merge_result_data + correctly interprets it as a single IMAGE output. + """ + preview_path = os.path.join( + output_dir, "previews", f"actor_{actor_index}_{photo_index}.jpg" + ) + if not os.path.isfile(preview_path): + print(f"[SelectActorPreview] File not found: {preview_path}") + return (torch.zeros(1, 512, 512, 3, dtype=torch.float32),) + + import cv2 + + img_bgr = cv2.imread(preview_path) + if img_bgr is None: + return (torch.zeros(1, 512, 512, 3, dtype=torch.float32),) + + img_rgb = img_bgr[:, :, ::-1].copy() + img_float = img_rgb.astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_float).unsqueeze(0) # [1, H, W, 3] + return (tensor,) + + +# ComfyUI node registration +NODE_CLASS_MAPPINGS = { + "SelectActorPreview": SelectActorPreview, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SelectActorPreview": "Select Actor Preview", +} diff --git a/pipeline/identity.py b/pipeline/identity.py index be8fe03..f3bcb91 100644 --- a/pipeline/identity.py +++ b/pipeline/identity.py @@ -11,7 +11,6 @@ import math import os import numpy as np from typing import Dict, List, Tuple, Optional -from collections import defaultdict from pipeline.tracker import FrameRecord from pipeline.detector import BoundingBox @@ -21,19 +20,18 @@ from core.config import DEFAULT_FACE_THRESHOLD, DEFAULT_MIN_FACE_CONFIDENCE class IdentityCluster: """Identity clustering using InsightFace face embeddings.""" - def __init__(self, threshold: float = DEFAULT_FACE_THRESHOLD): + def __init__(self, threshold: float = DEFAULT_FACE_THRESHOLD, model_dir: str = ""): self.threshold = threshold self.model = None self._loaded = False + self.model_dir = model_dir def _ensure_loaded(self): """Lazy-load InsightFace model.""" if self._loaded: return try: - import insightface from insightface.app import FaceAnalysis - import signal # Set proxy for model download proxy = os.environ.get("http_proxy", os.environ.get("HTTP_PROXY", "")) @@ -41,7 +39,25 @@ class IdentityCluster: os.environ.setdefault("HTTP_PROXY", proxy) os.environ.setdefault("HTTPS_PROXY", proxy) - model_dir = os.path.expanduser("~/.insightface/models") + # Resolve model directory: caller-specified > ComfyUI models dir > default + model_dir = self.model_dir + if not model_dir or not os.path.isdir(os.path.join(model_dir, "buffalo_l")): + # Try ComfyUI folder_paths "video-actor-extract" + try: + import folder_paths + + paths = folder_paths.get_folder_paths("video-actor-extract") + for p in paths: + candidate = os.path.join(p, "buffalo_l") + if os.path.isdir(candidate): + model_dir = p + break + except Exception: + pass + + if not model_dir or not os.path.isdir(os.path.join(model_dir, "buffalo_l")): + # Fallback to default InsightFace location + model_dir = os.path.expanduser("~/.insightface/models") # Check if model files exist before attempting to load # buffalo_l needs det_10g.onnx and w600k_r50.onnx @@ -285,8 +301,6 @@ class IdentityCluster: key=lambda tid: 0 if track_embeddings[tid] is not None else 1, ) - merge_log: List[Tuple[int, int, float]] = [] - # Margin for spatial tiebreaker: when top candidates' adjusted # similarity scores differ by ≤ this value, use spatial centroid # distance to break the tie. Prevents incorrect merges between diff --git a/tests/diagnose_tracker.py b/tests/diagnose_tracker.py new file mode 100644 index 0000000..62aee7a --- /dev/null +++ b/tests/diagnose_tracker.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +"""Diagnose MaskTracker frame assignments for frames 85-115.""" + +import os, sys, cv2, numpy as np + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from pipeline.segmenter import PersonSegmenter +from pipeline.mask_tracker import MaskTracker +from core.config import DEFAULT_MAX_LOST_FRAMES + +# Load video +cap = cv2.VideoCapture(os.path.join(PROJECT_ROOT, "tests", "sample1.mp4")) +all_frames = [] +while True: + ret, f = cap.read() + if not ret: + break + all_frames.append(f) # BGR +cap.release() +print(f"Loaded {len(all_frames)} frames") + +# Init +seg = PersonSegmenter( + model_path="/Users/cherokee/App/ComfyUI/models/video-actor-extract/yolov8n-seg.pt" +) +tracker = MaskTracker(max_lost_frames=DEFAULT_MAX_LOST_FRAMES, match_threshold_px=150.0) + +# Process frames 190-289 (Track 3 creation zone) +for fi in range(190, 290): + frame_bgr = all_frames[fi] + masks = seg.detect_masks(frame_bgr) + + # Compute mask info + mask_info = [] + for mid, mask in masks: + ys, xs = np.where(mask) + if len(ys) > 0: + cx, cy = float(xs.mean()), float(ys.mean()) + area = int(mask.sum()) + mask_info.append((mid, cx, cy, area)) + + # Snapshot before update + actors_before = {} + for aid, actor in tracker.get_active_actors().items(): + if not actor.closed: + actors_before[aid] = (actor.last_centroid, actor.last_frame_idx) + + tracker.update(fi, masks, frame_bgr, seg) + + # Snapshot after update + actors_after = {} + for aid, actor in tracker.get_active_actors().items(): + actors_after[aid] = ( + actor.last_centroid, + actor.last_frame_idx, + actor.closed, + len(actor.frames), + ) + + if fi >= 190: + print(f"\nFrame {fi}: {len(masks)} masks") + for mid, cx, cy, area in mask_info: + assigned = "UNMATCHED" + for aid, (cent, lfi, closed, cnt) in sorted(actors_after.items()): + if lfi == fi and not closed: + dist = ((cx - cent[0]) ** 2 + (cy - cent[1]) ** 2) ** 0.5 + assigned = f"-> actor_{aid} (dist={dist:.0f}px)" + break + print(f" mask_{mid}: ({cx:.0f},{cy:.0f}) area={area:>7d} {assigned}") + + for aid, (cent, lfi) in sorted(actors_before.items()): + still_active = aid in actors_after + matched_this_frame = still_active and actors_after[aid][1] == fi + if ( + not matched_this_frame + and not actors_after.get(aid, (None, None, True, 0))[2] + ): + print( + f" actor_{aid}: UNMATCHED (centroid=({cent[0]:.0f},{cent[1]:.0f}), last={lfi})" + ) + +# Final summary +print("\n=== FINAL TRACK SUMMARY ===") +final = tracker.finish() +for aid, actor in sorted(final.items()): + fidx = [f[0] for f in actor.frames] + print(f" actor_{aid}: frames {min(fidx)}-{max(fidx)}, count={len(fidx)}") diff --git a/tests/sample1.mp4 b/tests/sample1.mp4 new file mode 100644 index 0000000..9e4f50e Binary files /dev/null and b/tests/sample1.mp4 differ diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..cf8a6e9 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +"""Standalone pipeline test — runs full VideoActorExtract without ComfyUI. + +Reads all frames from a test video, converts to ComfyUI IMAGE tensor format, +instantiates VideoActorExtractor, and calls extract(). + +Usage: + cd /Users/cherokee/Project/ComfyUI-VideoActorExtract && python tests/test_pipeline.py +""" + +import os +import sys +import json +import time + +# --------------------------------------------------------------------------- +# Ensure project root is importable +# --------------------------------------------------------------------------- +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +# 1. Standard library — already imported above + +# 2. Third-party +import cv2 +import numpy as np +import torch + +# 3. Local project +from nodes.actor_extractor import VideoActorExtractor + + +# --------------------------------------------------------------------------- +# Expected ground truth for sample1.mp4 +# --------------------------------------------------------------------------- +EXPECTED = { + "num_actors": 3, + "description": ( + "1 man (appears at beginning and end with a gap), " + "2 children (appear together in same shots)" + ), +} + + +# --------------------------------------------------------------------------- +# Frame loading +# --------------------------------------------------------------------------- +def load_video_as_tensor(video_path: str) -> torch.Tensor: + """Read all frames from a video file and convert to ComfyUI IMAGE format. + + ComfyUI IMAGE format: shape [B, H, W, C], RGB, float32, values 0-1. + + Args: + video_path: Absolute or relative path to an MP4 video file. + + Returns: + torch.Tensor of shape [B, H, W, C]. + + Raises: + ValueError: If the video cannot be opened or has zero frames. + """ + print(f"[TestPipeline] Opening video: {video_path}") + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise ValueError(f"Cannot open video: {video_path}") + + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = cap.get(cv2.CAP_PROP_FPS) + + print( + f"[TestPipeline] Video info: {width}x{height}, " + f"{fps:.1f} fps, {total_frames} frames" + ) + + if total_frames <= 0: + cap.release() + raise ValueError(f"Video reports zero frames: {video_path}") + + # Pre-allocate tensor — much faster than repeated torch.from_numpy + frames = [] + idx = 0 + t0 = time.time() + + while True: + ret, frame_bgr = cap.read() + if not ret: + break + + # BGR uint8 -> RGB float32, normalize to 0-1 + frame_rgb = frame_bgr[:, :, ::-1].copy() # BGR -> RGB + frame_float = frame_rgb.astype(np.float32) / 255.0 + frames.append(frame_float) + + idx += 1 + if idx % 50 == 0 or idx == total_frames: + elapsed = time.time() - t0 + rate = idx / elapsed if elapsed > 0 else 0 + print( + f"[TestPipeline] Loaded {idx}/{total_frames} frames " + f"({elapsed:.1f}s, {rate:.0f} fps)" + ) + + cap.release() + + elapsed_total = time.time() - t0 + print( + f"[TestPipeline] Frame loading complete: {len(frames)} frames in {elapsed_total:.1f}s" + ) + + if not frames: + raise ValueError(f"No frames read from video: {video_path}") + + # Stack into [B, H, W, C] tensor + tensor = torch.from_numpy(np.stack(frames, axis=0)) + print(f"[TestPipeline] Tensor shape: {tensor.shape}, dtype: {tensor.dtype}") + return tensor + + +# --------------------------------------------------------------------------- +# Result printing +# --------------------------------------------------------------------------- +def print_results( + actor_json_str: str, + output_dir: str, + preview_tensor: torch.Tensor, +) -> dict: + """Parse and print detailed pipeline results. + + Args: + actor_json_str: JSON string returned by extract(). + output_dir: Output directory path. + preview_tensor: Preview image tensor. + + Returns: + Parsed JSON dict for further inspection. + """ + result = json.loads(actor_json_str) + + actors = result.get("actors", []) + video_info = result.get("video_info", {}) + + print("\n" + "=" * 70) + print("[TestPipeline] PIPELINE RESULTS") + print("=" * 70) + + # Video info + print(f"\n[Video Info]") + print( + f" Resolution: {video_info.get('width', '?')}x{video_info.get('height', '?')}" + ) + print(f" FPS: {video_info.get('fps', '?')}") + print(f" Total frames: {video_info.get('total_frames', '?')}") + print(f" Duration: {video_info.get('duration_sec', 0):.2f}s") + + # Actor summary + num_actors = len(actors) + print(f"\n[Actors] Found {num_actors} actor(s)") + + if not actors: + print(" (no actors detected)") + else: + for actor in actors: + actor_id = actor.get("actor_id", "?") + segments = actor.get("segments", []) + total_frames = actor.get("total_frames", 0) + segment_count = actor.get("segment_count", len(segments)) + + print(f"\n Actor '{actor_id}':") + print(f" Total output frames: {total_frames}") + print(f" Segments: {segment_count}") + + for i, seg in enumerate(segments): + start = seg.get("start_frame", "?") + end = seg.get("end_frame", "?") + count = seg.get("frame_count", "?") + print(f" Segment {i + 1}: frames {start}-{end} ({count} frames)") + + # Check if video file was created + video_file = os.path.join(output_dir, f"{actor_id}.mp4") + if os.path.exists(video_file): + size_mb = os.path.getsize(video_file) / (1024 * 1024) + print(f" Video: {video_file} ({size_mb:.2f} MB)") + else: + print(f" Video: NOT FOUND at {video_file}") + + # Preview tensor + if preview_tensor is not None and preview_tensor.numel() > 0: + print(f"\n[Preview] Tensor shape: {preview_tensor.shape}") + else: + print(f"\n[Preview] No preview frames generated") + + # Output directory + print(f"\n[Output] Directory: {output_dir}") + json_file = os.path.join(output_dir, "actor_info.json") + if os.path.exists(json_file): + print(f" JSON: {json_file}") + else: + print(f" JSON: NOT FOUND at {json_file}") + + print("=" * 70) + + return result + + +# --------------------------------------------------------------------------- +# Expected vs Actual comparison +# --------------------------------------------------------------------------- +def print_comparison(result: dict) -> None: + """Compare actual pipeline results against expected ground truth. + + Args: + result: Parsed JSON dict from the pipeline. + """ + actors = result.get("actors", []) + actual_num = len(actors) + + print("\n" + "=" * 70) + print("[TestPipeline] EXPECTED vs ACTUAL") + print("=" * 70) + + print(f"\n Expected actors: {EXPECTED['num_actors']}") + print(f" Actual actors: {actual_num}") + print( + f" Match: {'PASS' if actual_num == EXPECTED['num_actors'] else 'MISMATCH'}" + ) + print(f"\n Expected description: {EXPECTED['description']}") + + # Per-actor segment breakdown for diagnosis + if actual_num != EXPECTED["num_actors"]: + print(f"\n [Diagnosis]") + if actual_num > EXPECTED["num_actors"]: + print( + f" More actors detected than expected. " + f"Possible causes: split tracks, false positives, or identity clustering too strict." + ) + else: + print( + f" Fewer actors detected than expected. " + f"Possible causes: identity clustering too aggressive (merging distinct people), " + f"or short tracks filtered out by min_track_length." + ) + + # Print detailed per-actor frame ranges for manual inspection + print(f"\n [Per-actor frame ranges for manual inspection]") + for actor in actors: + actor_id = actor.get("actor_id", "?") + segments = actor.get("segments", []) + ranges = ", ".join(f"{s['start_frame']}-{s['end_frame']}" for s in segments) + print(f" Actor '{actor_id}': [{ranges}]") + + print("=" * 70) + + +# --------------------------------------------------------------------------- +# Save results to JSON +# --------------------------------------------------------------------------- +def save_results( + result: dict, + output_dir: str, + save_path: str, +) -> None: + """Save pipeline results plus metadata to a JSON file for inspection. + + Args: + result: Parsed JSON dict from the pipeline. + output_dir: Pipeline output directory path. + save_path: Where to write the combined results JSON. + """ + report = { + "expected": EXPECTED, + "actual_num_actors": len(result.get("actors", {})), + "match": len(result.get("actors", {})) == EXPECTED["num_actors"], + "output_dir": output_dir, + "pipeline_result": result, + } + + os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) + with open(save_path, "w") as f: + json.dump(report, f, indent=2, default=str) + + print(f"[TestPipeline] Results saved to: {save_path}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main() -> None: + """Run the full pipeline test on sample1.mp4.""" + print("[TestPipeline] Starting standalone pipeline test") + print(f"[TestPipeline] Project root: {PROJECT_ROOT}") + + # Resolve video path + video_path = os.path.join(PROJECT_ROOT, "tests", "sample1.mp4") + if not os.path.exists(video_path): + print(f"[TestPipeline] ERROR: Test video not found: {video_path}") + sys.exit(1) + + # Results save path + results_save_path = os.path.join( + PROJECT_ROOT, "tests", "test_pipeline_results.json" + ) + + # Step 1: Load video frames into ComfyUI IMAGE tensor + print("\n[TestPipeline] Step 1: Loading video frames...") + images = load_video_as_tensor(video_path) + + # Step 2: Instantiate the extractor + print("\n[TestPipeline] Step 2: Instantiating VideoActorExtractor...") + extractor = VideoActorExtractor() + + # Step 3: Run the full pipeline + print("\n[TestPipeline] Step 3: Running extract() pipeline...") + print("-" * 70) + t0 = time.time() + + actor_json_str, output_dir, preview_tensor = extractor.extract( + images=images, + video_path=video_path, + ) + + elapsed = time.time() - t0 + print("-" * 70) + print(f"[TestPipeline] Pipeline completed in {elapsed:.1f}s") + + # Step 4: Print detailed results + result = print_results(actor_json_str, output_dir, preview_tensor) + + # Step 5: Compare with expected ground truth + print_comparison(result) + + # Step 6: Save results JSON + save_results(result, output_dir, results_save_path) + + # Final verdict + actual_num = len(result.get("actors", {})) + if actual_num == EXPECTED["num_actors"]: + print(f"\n[TestPipeline] RESULT: PASS ({actual_num} actors detected)") + else: + print( + f"\n[TestPipeline] RESULT: MISMATCH " + f"(expected {EXPECTED['num_actors']}, got {actual_num})" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_pipeline_results.json b/tests/test_pipeline_results.json new file mode 100644 index 0000000..d54f117 --- /dev/null +++ b/tests/test_pipeline_results.json @@ -0,0 +1,73 @@ +{ + "expected": { + "num_actors": 3, + "description": "1 man (appears at beginning and end with a gap), 2 children (appear together in same shots)" + }, + "actual_num_actors": 3, + "match": true, + "output_dir": "/Users/cherokee/output/ComfyUI-VideoActorExtract/1465a0a8", + "pipeline_result": { + "actors": [ + { + "actor_id": "actor_0", + "segment_count": 2, + "total_frames": 125, + "segments": [ + { + "segment_id": 0, + "start_frame": 0, + "end_frame": 96, + "start_time_sec": 0.0, + "end_time_sec": 4.0, + "frame_count": 97 + }, + { + "segment_id": 1, + "start_frame": 261, + "end_frame": 288, + "start_time_sec": 10.88, + "end_time_sec": 12.0, + "frame_count": 28 + } + ] + }, + { + "actor_id": "actor_1", + "segment_count": 1, + "total_frames": 100, + "segments": [ + { + "segment_id": 0, + "start_frame": 97, + "end_frame": 196, + "start_time_sec": 4.04, + "end_time_sec": 8.17, + "frame_count": 100 + } + ] + }, + { + "actor_id": "actor_2", + "segment_count": 1, + "total_frames": 164, + "segments": [ + { + "segment_id": 0, + "start_frame": 97, + "end_frame": 260, + "start_time_sec": 4.04, + "end_time_sec": 10.83, + "frame_count": 164 + } + ] + } + ], + "video_info": { + "total_frames": 290, + "fps": 24.0, + "width": 1080, + "height": 1916, + "duration_sec": 12.083333333333334 + } + } +} \ No newline at end of file