feat: add SelectActorPreview node, identity temporal constraints, and preview images to disk
- Add SelectActorPreview node to load per-actor preview images as ComfyUI IMAGE tensors
- Fix SelectActorPreview return value: must return (tensor,) tuple, not bare tensor,
to prevent ComfyUI's merge_result_data from misinterpreting batch dimension as
output count (root cause of 1920-image bug)
- Change VideoActorExtractor third output from IMAGE tensor to INT actor_count
- Save top-5 preview JPGs per actor to {output_dir}/previews/ instead of
building a large preview tensor in memory
- Unify InsightFace model path to ComfyUI/models/video-actor-extract/buffalo_l/
- Add temporal gap penalty with grace gap to identity clustering
- Add temporal boost with Gaussian decay and spatial blocker check
- Add dual-signal track splitting (centroid displacement + bbox area ratio)
- Add spatial centroid proximity tiebreaker for identity merging
- Add debris mask filtering (min_mask_area=20000) in MaskTracker
- Clean up unused imports across modified files
- Add AGENTS.md project guidelines and test infrastructure
This commit is contained in:
Cherokee
@@ -8,3 +8,4 @@ build/
|
||||
.eggs/
|
||||
test_output/
|
||||
output/
|
||||
models/
|
||||
|
||||
@@ -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.
|
||||
+38
-14
@@ -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"]
|
||||
|
||||
+38
-48
@@ -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
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
+21
-7
@@ -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
|
||||
|
||||
@@ -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)}")
|
||||
Binary file not shown.
@@ -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()
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user