diff --git a/core/config.py b/core/config.py index f55c4fc..481e265 100644 --- a/core/config.py +++ b/core/config.py @@ -14,7 +14,7 @@ DEFAULT_MAX_LOST_FRAMES = 30 # Face recognition DEFAULT_FACE_THRESHOLD = 0.6 -DEFAULT_MIN_FACE_CONFIDENCE = 0.5 +DEFAULT_MIN_FACE_CONFIDENCE = 0.3 # Video processing DEFAULT_FPS_SAMPLE = 3 diff --git a/nodes/actor_extractor.py b/nodes/actor_extractor.py index cfb357a..4a6e4e5 100644 --- a/nodes/actor_extractor.py +++ b/nodes/actor_extractor.py @@ -76,6 +76,41 @@ def _get_comfyui_output_dir() -> str: return output_dir +def _resolve_model_path(name: str) -> str: + """Resolve a model filename to an absolute path. + + Resolution order: + 1. If already an absolute path and file exists, return as-is. + 2. Try ComfyUI folder_paths "video-actor-extract" folder. + 3. Try ~/App/ComfyUI/models/video-actor-extract/. + 4. Return the bare name (backward compat / let ultralytics download it). + """ + # Already absolute and exists? + if os.path.isabs(name) and os.path.isfile(name): + return name + + # 1. Try ComfyUI folder_paths + try: + import folder_paths + + paths = folder_paths.get_folder_paths("video-actor-extract") + for p in paths: + candidate = os.path.join(p, name) + if os.path.isfile(candidate): + return candidate + except Exception: + pass + + # 2. Try default ComfyUI models directory + default_dir = os.path.expanduser("~/App/ComfyUI/models/video-actor-extract/") + candidate = os.path.join(default_dir, name) + if os.path.isfile(candidate): + return candidate + + # 3. Return bare name (backward compat) + return name + + def _build_continuous_segments( actor_frames: List[Tuple[int, np.ndarray, int]], max_gap: int = 30, @@ -234,7 +269,7 @@ class VideoActorExtractor: "seg_model_path": ( "STRING", { - "default": DEFAULT_SEG_MODEL, + "default": "yolov8n-seg.pt", "multiline": False, "tooltip": "YOLOv8-seg model for person segmentation", }, @@ -310,6 +345,12 @@ class VideoActorExtractor: output_dir = _get_comfyui_output_dir() print(f"[VideoActorExtract] Output directory: {output_dir}") + # Resolve model paths + model_path = _resolve_model_path(model_path) + seg_model_path = _resolve_model_path(seg_model_path) + print(f"[VideoActorExtract] Resolved model_path: {model_path}") + print(f"[VideoActorExtract] Resolved seg_model_path: {seg_model_path}") + # ---------------------------------------------------------------- # Convert IMAGE tensor to numpy frames (RGB 0-1 -> BGR 0-255) # ---------------------------------------------------------------- diff --git a/pipeline/identity.py b/pipeline/identity.py index 6afbcc3..b71cc0e 100644 --- a/pipeline/identity.py +++ b/pipeline/identity.py @@ -16,12 +16,12 @@ 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): self.threshold = threshold self.model = None self._loaded = False - + def _ensure_loaded(self): """Lazy-load InsightFace model.""" if self._loaded: @@ -30,29 +30,42 @@ class IdentityCluster: 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', '')) + 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) - + 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'] + 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) - + 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] 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"], @@ -66,94 +79,81 @@ class IdentityCluster: print("[Identity] Will use track-only mode (no face clustering)") self._loaded = False self.model = None - - def _get_face_embedding(self, frame: np.ndarray, bbox: BoundingBox) -> Optional[np.ndarray]: + + def _get_face_embedding( + self, frame: np.ndarray, bbox: BoundingBox + ) -> Optional[np.ndarray]: """ Extract face embedding from a cropped region of a frame. - + Args: frame: Full BGR frame bbox: Bounding box of the person - + Returns: Face embedding vector or None if no face detected """ self._ensure_loaded() if self.model is None: return None - - # Crop person region with some padding + + # Crop person region — expand bbox by 40% to ensure full face is included h, w = frame.shape[:2] - pad_x = int(bbox.width * 0.2) - pad_y = int(bbox.height * 0.3) - x1 = max(0, int(bbox.x1) - pad_x) - y1 = max(0, int(bbox.y1) - pad_y) - x2 = min(w, int(bbox.x2) + pad_x) - y2 = min(h, int(bbox.y2) + pad_y) - + expand = 0.4 + bw = bbox.width * expand + bh = bbox.height * expand + x1 = max(0, int(bbox.x1 - bw)) + y1 = max(0, int(bbox.y1 - bh)) + x2 = min(w, int(bbox.x2 + bw)) + y2 = min(h, int(bbox.y2 + bh)) + crop = frame[y1:y2, x1:x2] if crop.size == 0: return None - + faces = self.model.get(crop) if not faces: return None - + # Return the embedding of the face with highest detection confidence best_face = max(faces, key=lambda f: f.det_score) if best_face.det_score < DEFAULT_MIN_FACE_CONFIDENCE: return None - + return best_face.embedding - - def _get_track_embedding(self, track_records: List[FrameRecord], frames: np.ndarray) -> Optional[np.ndarray]: - """ - Get aggregated face embedding for a track by sampling key frames. - - Args: - track_records: List of FrameRecord for this track - frames: All sampled frames (indexed by original frame index) - - Returns: - Average face embedding or None - """ + + def _get_track_embedding(self, track_records, frames): + """Get aggregated face embedding for a track by trying multiple frames.""" embeddings = [] - - # Sample up to 5 key frames from the track n = len(track_records) if n == 0: return None - - if n <= 5: + + # Try up to 10 evenly-spaced frames to maximize face detection chances + if n <= 10: indices = list(range(n)) else: - step = n // 5 - indices = list(range(0, n, step))[:5] - + step = n / 10.0 + indices = [int(i * step) for i in range(10)] + for idx in indices: rec = track_records[idx] - # Map frame_idx to index in frames array - # frames is indexed by the sampling, so we need the original frame mapping - # We'll store the frame array alongside the records - # For now, assume frames are accessible by frame_idx - # This will be handled by the caller frame = frames.get(rec.frame_idx, None) if frame is None: continue - bbox = BoundingBox(rec.x1, rec.y1, rec.x2, rec.y2) emb = self._get_face_embedding(frame, bbox) if emb is not None: embeddings.append(emb) - + if not embeddings: return None - + # Average embeddings and normalize avg_emb = np.mean(embeddings, axis=0) avg_emb = avg_emb / (np.linalg.norm(avg_emb) + 1e-8) return avg_emb - + def cluster_tracks( self, track_records: Dict[int, List[FrameRecord]], @@ -162,48 +162,49 @@ class IdentityCluster: ) -> Dict[int, str]: """ Cluster tracks into actor identities. - + Args: track_records: {track_id: [FrameRecord, ...]} frames: {frame_idx: numpy_frame} - all frames needed for face extraction min_track_length: Minimum number of frames for a track to be considered - + Returns: {track_id: "actor_0", ...} """ self._ensure_loaded() - + # Filter out short tracks valid_tracks = { - tid: recs for tid, recs in track_records.items() + tid: recs + for tid, recs in track_records.items() if len(recs) >= min_track_length } - + if not valid_tracks: return {} - + # If no face model, assign each track a unique actor ID if self.model is None: return {tid: f"actor_{i}" for i, tid in enumerate(valid_tracks.keys())} - + # Extract face embeddings for each track track_embeddings: Dict[int, Optional[np.ndarray]] = {} for tid, recs in valid_tracks.items(): emb = self._get_track_embedding(recs, frames) track_embeddings[tid] = emb - + # Greedy clustering by cosine similarity actor_id_counter = 0 track_to_actor: Dict[int, str] = {} # actor_id -> list of embeddings belonging to this actor actor_embeddings: Dict[str, List[np.ndarray]] = {} - + # Sort tracks by embedding quality (non-None first) sorted_tracks = sorted( valid_tracks.keys(), key=lambda tid: 0 if track_embeddings[tid] is not None else 1, ) - + for tid in sorted_tracks: emb = track_embeddings[tid] if emb is None: @@ -212,11 +213,11 @@ class IdentityCluster: actor_id_counter += 1 track_to_actor[tid] = actor_id continue - + # Find best matching existing actor best_similarity = -1 best_actor = None - + for actor_id, embs in actor_embeddings.items(): # Average similarity to all embeddings in this actor group sims = [float(np.dot(emb, e)) for e in embs] @@ -224,7 +225,7 @@ class IdentityCluster: if avg_sim > best_similarity: best_similarity = avg_sim best_actor = actor_id - + if best_similarity >= self.threshold: # Merge into existing actor track_to_actor[tid] = best_actor @@ -235,5 +236,117 @@ class IdentityCluster: actor_id_counter += 1 track_to_actor[tid] = actor_id actor_embeddings[actor_id] = [emb] - + + # Spatial-temporal fallback: merge faceless tracks into existing actors + # by checking time-range overlap and bbox overlap + faceless_tracks = [tid for tid, emb in track_embeddings.items() if emb is None] + if faceless_tracks: + self._merge_faceless_by_overlap( + faceless_tracks, + valid_tracks, + track_to_actor, + actor_embeddings, + actor_id_counter, + ) + # Renumber actor IDs sequentially + actor_map = { + aid: idx for idx, aid in enumerate(sorted(set(track_to_actor.values()))) + } + track_to_actor = {tid: f"actor_{v}" for tid, v in actor_map.items()} + return track_to_actor + + def _merge_faceless_by_overlap( + self, + faceless_tracks: list, + valid_tracks: dict, + track_to_actor: dict, + actor_embeddings: dict, + actor_id_counter: int, + ): + """Try to assign faceless tracks to existing actors by spatial-temporal overlap.""" + # Build frame-range per existing actor + actor_frame_ranges: Dict[str, Tuple[int, int]] = {} + actor_bboxes: Dict[str, List[BoundingBox]] = {} + for tid, aid in track_to_actor.items(): + if aid not in actor_frame_ranges: + actor_frame_ranges[aid] = (float("inf"), float("-inf")) + actor_bboxes[aid] = [] + recs = valid_tracks[tid] + frame_indices = [r.frame_idx for r in recs] + lo, hi = actor_frame_ranges[aid] + actor_frame_ranges[aid] = ( + min(lo, min(frame_indices)), + max(hi, max(frame_indices)), + ) + for r in recs: + actor_bboxes[aid].append(BoundingBox(r.x1, r.y1, r.x2, r.y2)) + + for tid in faceless_tracks: + recs = valid_tracks[tid] + track_frames = set(r.frame_idx for r in recs) + track_lo, track_hi = min(track_frames), max(track_frames) + + best_actor = None + best_score = 0.0 + + for aid, (a_lo, a_hi) in actor_frame_ranges.items(): + # Time range overlap + overlap_start = max(track_lo, a_lo) + overlap_end = min(track_hi, a_hi) + if overlap_start > overlap_end: + continue + time_overlap = (overlap_end - overlap_start + 1) / max( + 1, track_hi - track_lo + 1 + ) + if time_overlap < 0.3: + continue + + # Bbox overlap (use median bbox of each) + track_bboxes = [BoundingBox(r.x1, r.y1, r.x2, r.y2) for r in recs] + med_track = self._median_bbox(track_bboxes) + med_actor = self._median_bbox(actor_bboxes[aid]) + iou = self._bbox_iou(med_track, med_actor) + + score = time_overlap * 0.5 + iou * 0.5 + if score > best_score: + best_score = score + best_actor = aid + + if best_actor is not None and best_score > 0.3: + track_to_actor[tid] = best_actor + print( + f"[Identity] Merged faceless track {tid} -> {best_actor} (score={best_score:.2f})" + ) + else: + # Assign a new unique ID + new_id = f"actor_{actor_id_counter}" + actor_id_counter += 1 + track_to_actor[tid] = new_id + print(f"[Identity] Faceless track {tid} unmatched, assigned {new_id}") + + @staticmethod + def _median_bbox(bboxes): + """Compute median bounding box from a list.""" + if not bboxes: + return None + xs = [b.x1 for b in bboxes] + [b.x2 for b in bboxes] + ys = [b.y1 for b in bboxes] + [b.y2 for b in bboxes] + mid_x = sorted(xs)[len(xs) // 2] + mid_y = sorted(ys)[len(ys) // 2] + return BoundingBox(mid_x - 1, mid_y - 1, mid_x + 1, mid_y + 1) + + @staticmethod + def _bbox_iou(a, b): + """Compute IoU between two BoundingBoxes.""" + if a is None or b is None: + return 0.0 + ix1 = max(a.x1, b.x1) + iy1 = max(a.y1, b.y1) + ix2 = min(a.x2, b.x2) + iy2 = min(a.y2, b.y2) + inter = max(0, ix2 - ix1) * max(0, iy2 - iy1) + area_a = max(0, a.x2 - a.x1) * max(0, a.y2 - a.y1) + area_b = max(0, b.x2 - b.x1) * max(0, b.y2 - b.y1) + union = area_a + area_b - inter + return inter / (union + 1e-8)