diff --git a/core/__pycache__/__init__.cpython-310.pyc b/core/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 3628d4a..0000000 Binary files a/core/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/core/__pycache__/config.cpython-310.pyc b/core/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 0bd3b60..0000000 Binary files a/core/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/core/__pycache__/video_reader.cpython-310.pyc b/core/__pycache__/video_reader.cpython-310.pyc deleted file mode 100644 index 8aade92..0000000 Binary files a/core/__pycache__/video_reader.cpython-310.pyc and /dev/null differ diff --git a/nodes/__init__.py b/nodes/__init__.py index b24bf7a..9014f8c 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -5,6 +5,10 @@ Exports the main VideoActorExtractor node for use in ComfyUI workflows. from __future__ import annotations -from .actor_extractor import VideoActorExtractor +from .actor_extractor import VideoActorExtractor, NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ["VideoActorExtractor"] + +__all__.extend(NODE_CLASS_MAPPINGS.keys()) + +WEB_DIRECTORY = "./js" diff --git a/nodes/__pycache__/__init__.cpython-310.pyc b/nodes/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 86eec32..0000000 Binary files a/nodes/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/nodes/__pycache__/actor_extractor.cpython-310.pyc b/nodes/__pycache__/actor_extractor.cpython-310.pyc deleted file mode 100644 index 50fc681..0000000 Binary files a/nodes/__pycache__/actor_extractor.cpython-310.pyc and /dev/null differ diff --git a/pipeline/__pycache__/__init__.cpython-310.pyc b/pipeline/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 2d924d4..0000000 Binary files a/pipeline/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/pipeline/__pycache__/cropper.cpython-310.pyc b/pipeline/__pycache__/cropper.cpython-310.pyc deleted file mode 100644 index 320719f..0000000 Binary files a/pipeline/__pycache__/cropper.cpython-310.pyc and /dev/null differ diff --git a/pipeline/__pycache__/detector.cpython-310.pyc b/pipeline/__pycache__/detector.cpython-310.pyc deleted file mode 100644 index 3eefe25..0000000 Binary files a/pipeline/__pycache__/detector.cpython-310.pyc and /dev/null differ diff --git a/pipeline/__pycache__/identity.cpython-310.pyc b/pipeline/__pycache__/identity.cpython-310.pyc deleted file mode 100644 index 42cc869..0000000 Binary files a/pipeline/__pycache__/identity.cpython-310.pyc and /dev/null differ diff --git a/pipeline/__pycache__/merger.cpython-310.pyc b/pipeline/__pycache__/merger.cpython-310.pyc deleted file mode 100644 index 6c91258..0000000 Binary files a/pipeline/__pycache__/merger.cpython-310.pyc and /dev/null differ diff --git a/pipeline/__pycache__/tracker.cpython-310.pyc b/pipeline/__pycache__/tracker.cpython-310.pyc deleted file mode 100644 index f066a68..0000000 Binary files a/pipeline/__pycache__/tracker.cpython-310.pyc and /dev/null differ diff --git a/pipeline/detector.py b/pipeline/detector.py index 518576f..b8a95fd 100644 --- a/pipeline/detector.py +++ b/pipeline/detector.py @@ -119,39 +119,20 @@ class PersonDetector: def detect_batch(self, frames: np.ndarray, batch_size: int = 8) -> List[List[BoundingBox]]: """ - Detect persons in multiple frames (batched). + Detect persons in multiple frames. Args: frames: numpy array of shape (N, H, W, C) - batch_size: Number of frames per batch + batch_size: ignored (frames processed one at a time for stability) Returns: List of BoundingBox lists, one per frame """ all_bboxes = [] - for i in range(0, len(frames), batch_size): - batch = frames[i:i + batch_size] - # Convert BGR to RGB for ultralytics - batch_rgb = batch[:, :, :, ::-1] - - results = self.model( - batch_rgb, - classes=[0], - conf=self.conf_threshold, - iou=self.iou_threshold, - verbose=False, - device=self.device, - ) - - for result in results: - bboxes = [] - boxes = result.boxes - if boxes is not None: - for box in boxes: - x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() - conf = float(box.conf[0].cpu().numpy()) - bboxes.append(BoundingBox(x1, y1, x2, y2, conf)) - all_bboxes.append(bboxes) + for i in range(len(frames)): + frame = frames[i] + bboxes = self.detect(frame) + all_bboxes.append(bboxes) return all_bboxes diff --git a/pipeline/identity.py b/pipeline/identity.py index 091143f..6afbcc3 100644 --- a/pipeline/identity.py +++ b/pipeline/identity.py @@ -4,6 +4,7 @@ Uses InsightFace for face embedding extraction and cosine similarity clustering to merge tracks belonging to the same actor. """ +import os import numpy as np from typing import Dict, List, Tuple, Optional from collections import defaultdict @@ -28,12 +29,36 @@ class IdentityCluster: 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', '')) + if proxy: + os.environ.setdefault('HTTP_PROXY', proxy) + os.environ.setdefault('HTTPS_PROXY', proxy) + + 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 + required_files = ['1k3d68.onnx', '2d106det.onnx', 'det_10g.onnx', 'genderage.onnx', 'w600k_mbf.onnx', 'w600k_r50.onnx'] + model_path = os.path.join(model_dir, "buffalo_l") + has_all = all(os.path.exists(os.path.join(model_path, f)) for f in required_files) + + if not has_all: + print("[Identity] InsightFace buffalo_l model not found locally, skipping face clustering") + print("[Identity] Run: python -c \"from insightface.app import FaceAnalysis; FaceAnalysis(name='buffalo_l').prepare(ctx_id=-1)\"") + print("[Identity] to download the model (~350MB)") + self._loaded = False + self.model = None + return self.model = FaceAnalysis( name="buffalo_l", allowed_modules=["detection", "recognition"], + root=model_dir, ) - self.model.prepare(ctx_id=0, det_size=(640, 640)) + self.model.prepare(ctx_id=-1, det_size=(640, 640)) # ctx_id=-1 for CPU self._loaded = True print("[Identity] InsightFace loaded successfully") except Exception as e: