""" ShortSmith v4 - Music Highlight Extractor Gradio UI for extracting the best 15-second highlights from music videos and live performances. """ import gradio as gr import json import os import shutil import tempfile import urllib.request import zipfile from models.object_detector import ObjectDetector from pathlib import Path from typing import Optional, List, Dict, Any, Tuple from core.video_processor import VideoProcessor from pipeline.orchestrator import MusicHighlightOrchestrator, PipelineResult from utils.logger import get_logger import config as cfg logger = get_logger("app") def _insightface_model_url() -> str: return os.environ.get( "INSIGHTFACE_MODEL_URL", "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip", ) def _insightface_model_dir() -> Path: # Allow overriding cache root (useful for HF Spaces persistent storage, e.g. /data/.insightface) root = Path(os.environ.get("INSIGHTFACE_HOME", str(Path.home() / ".insightface"))) return root / "models" / "buffalo_l" def _insightface_required_files() -> set[str]: return { "1k3d68.onnx", "2d106det.onnx", "det_10g.onnx", "genderage.onnx", "w600k_r50.onnx", } def get_insightface_weight_status() -> dict: """Return current local model status.""" model_dir = _insightface_model_dir() required = _insightface_required_files() present = {p.name for p in model_dir.glob("*.onnx")} if model_dir.exists() else set() missing = sorted(required - present) return { "installed": len(missing) == 0, "model_dir": str(model_dir), "source_url": _insightface_model_url(), "present": sorted(present), "missing": missing, } def ensure_insightface_weights() -> None: """Ensure InsightFace buffalo_l weights exist in local cache.""" model_dir = _insightface_model_dir() required = _insightface_required_files() if model_dir.exists() and required.issubset({p.name for p in model_dir.glob("*.onnx")}): logger.info("InsightFace weights already available locally") return model_dir.mkdir(parents=True, exist_ok=True) url = _insightface_model_url() logger.info(f"Downloading InsightFace model pack from: {url}") with tempfile.TemporaryDirectory(prefix="insightface_dl_") as td: zip_path = Path(td) / "buffalo_l.zip" extract_dir = Path(td) / "extract" extract_dir.mkdir(parents=True, exist_ok=True) # Use an explicit User-Agent to avoid blocked anonymous requests in some environments. req = urllib.request.Request(url, headers={"User-Agent": "ShortSmith/1.0"}) with urllib.request.urlopen(req, timeout=180) as resp, open(zip_path, "wb") as out: out.write(resp.read()) with zipfile.ZipFile(zip_path, "r") as zf: zf.extractall(extract_dir) # Release archives can contain files at root or under a folder; support both. extracted = list(extract_dir.rglob("*.onnx")) copied = 0 for src in extracted: if src.name in required: shutil.copy2(src, model_dir / src.name) copied += 1 final_files = {p.name for p in model_dir.glob("*.onnx")} missing = required - final_files if missing: raise RuntimeError(f"InsightFace model setup incomplete; missing: {sorted(missing)}") logger.info(f"InsightFace weights ready in {model_dir}") # Singleton orchestrator — keeps Whisper model loaded between runs _orchestrator: Optional[MusicHighlightOrchestrator] = None def get_orchestrator(progress_callback=None) -> MusicHighlightOrchestrator: global _orchestrator if _orchestrator is None: _orchestrator = MusicHighlightOrchestrator(progress_callback=progress_callback) else: _orchestrator.progress_callback = progress_callback return _orchestrator def format_timestamp(seconds: float) -> str: m = int(seconds) // 60 s = int(seconds) % 60 return f"{m}:{s:02d}" def _empty_face_clip_outputs(status: str = "") -> tuple: return ( None, None, None, None, None, "", "", "", "", "", status, ) def _get_video_duration(video_path: str) -> float: try: import cv2 cap = cv2.VideoCapture(video_path) fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) cap.release() if fps > 0 and frame_count > 0: return float(frame_count / fps) except Exception: pass return 0.0 def _build_face_segments( occurrences: List[Dict[str, Any]], clip_duration: float, video_duration: float, effective_rate: float, ) -> List[Tuple[float, float]]: """Create clip segments for all face occurrences (fixed-duration or auto ranges).""" if not occurrences: return [] # Deduplicate by frame to avoid repeated timestamps. unique: List[Dict[str, Any]] = [] seen_frames = set() for occ in sorted(occurrences, key=lambda x: float(x.get("timestamp", 0.0))): frame_index = int(occ.get("frame_index", -1)) if frame_index in seen_frames: continue seen_frames.add(frame_index) unique.append(occ) if not unique: return [] segs: List[Tuple[float, float]] = [] def clamp_segment(start: float, end: float) -> Optional[Tuple[float, float]]: s = max(0.0, float(start)) e = max(s, float(end)) if video_duration > 0: if e > video_duration: shift = e - video_duration s = max(0.0, s - shift) e = video_duration if e - s <= 0.15: return None return (s, e) if clip_duration > 0: duration = float(clip_duration) # For fixed duration, build windows for all timestamps then merge overlaps. windows: List[Tuple[float, float]] = [] for occ in unique: ts = float(occ.get("timestamp", 0.0)) half = duration / 2.0 seg = clamp_segment(ts - half, ts + half) if seg is not None: windows.append(seg) if not windows: return [] windows.sort(key=lambda x: x[0]) merged: List[Tuple[float, float]] = [windows[0]] merge_gap = max(0.15, duration * 0.35) for s, e in windows[1:]: ls, le = merged[-1] if s <= le + merge_gap: merged[-1] = (ls, max(le, e)) else: merged.append((s, e)) return merged # Auto mode (clip_duration == 0): group contiguous visibility windows. step = 1.0 / max(0.1, float(effective_rate)) gap_threshold = max(0.6, step * 3.0) ts_list = [float(occ.get("timestamp", 0.0)) for occ in unique] groups: List[Tuple[float, float]] = [] start = ts_list[0] prev = ts_list[0] for ts in ts_list[1:]: if ts - prev <= gap_threshold: prev = ts continue groups.append((start, prev)) start = ts prev = ts groups.append((start, prev)) for g_start, g_end in groups: seg_start = g_start - 0.2 seg_end = g_end + 0.2 if seg_end - seg_start < 1.0: center = (g_start + g_end) / 2.0 seg_start = center - 0.5 seg_end = center + 0.5 seg = clamp_segment(seg_start, seg_end) if seg is not None: segs.append(seg) return segs def _cosine_similarity(v1: List[float], v2: List[float]) -> float: if not v1 or not v2 or len(v1) != len(v2): return -1.0 dot = 0.0 n1 = 0.0 n2 = 0.0 for a, b in zip(v1, v2): fa = float(a) fb = float(b) dot += fa * fb n1 += fa * fa n2 += fb * fb if n1 <= 0.0 or n2 <= 0.0: return -1.0 return float(dot / ((n1 ** 0.5) * (n2 ** 0.5))) def _collect_matched_identities( identities: List[Dict[str, Any]], selected_identity_id: int, similarity_threshold: float = 0.45, ) -> List[Dict[str, Any]]: """Match selected identity to similar embedding-centroids and return all matched identities.""" by_id: Dict[int, Dict[str, Any]] = {} for ident in identities: try: by_id[int(ident.get("identity_id"))] = ident except Exception: continue if selected_identity_id not in by_id: return [] selected = by_id[selected_identity_id] selected_emb = selected.get("embedding_mean") if not selected_emb: return [selected] matched: List[Dict[str, Any]] = [] for ident in by_id.values(): emb = ident.get("embedding_mean") sim = _cosine_similarity(selected_emb, emb) if emb else -1.0 if sim >= similarity_threshold or int(ident.get("identity_id")) == selected_identity_id: ident_copy = dict(ident) ident_copy["embedding_similarity_to_selected"] = sim matched.append(ident_copy) matched.sort(key=lambda x: float(x.get("embedding_similarity_to_selected", -1.0)), reverse=True) return matched def process_video( video_path: str, num_clips: int, progress: gr.Progress = gr.Progress(), ) -> tuple: """Main processing function called by Gradio.""" if not video_path: return ( None, None, None, "", "", "Please upload a video first." ) def on_progress(fraction: float, message: str): progress(fraction, desc=message) orch = get_orchestrator(progress_callback=on_progress) result: PipelineResult = orch.process( video_path=video_path, num_clips=num_clips, ) if not result.success: return ( None, None, None, "", "", f"Error: {result.error_message}" ) # Build clip outputs (up to 3 slots) clip_videos = [None, None, None] clip_labels = ["", "", ""] for i, clip in enumerate(result.clips[:3]): clip_videos[i] = str(clip.clip_path) start = format_timestamp(clip.start_time) end = format_timestamp(clip.end_time) clip_labels[i] = ( f"Clip {clip.rank} • {start} – {end} • Score: {clip.hype_score:.2f}" ) status = ( f"Done in {result.processing_time:.0f}s • " f"{len(result.clips)} clip(s) extracted" ) transcript = result.transcript or "(no lyrics detected)" return ( clip_videos[0], clip_videos[1], clip_videos[2], clip_labels[0], clip_labels[1], clip_labels[2], transcript, status, ) # --------------------------------------------------------------------------- # Build UI # --------------------------------------------------------------------------- with gr.Blocks(title="ShortSmith v4 — Music Highlights", theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 🎵 ShortSmith v4 — Music Highlight Extractor Upload a music video or live performance (≤ 15 min) and get the best **15-second highlights** automatically. **How it works:** Whisper transcribes the audio → chorus sections are detected → Librosa picks the highest-energy moments → FFmpeg cuts the clips. """ ) detection_state = gr.State(value={}) selected_identity_state = gr.State(value=None) with gr.Row(): with gr.Column(scale=1): video_input = gr.Video(label="Upload Music Video") num_clips_slider = gr.Slider( minimum=1, maximum=3, value=1, step=1, label="Number of clips to extract", ) run_btn = gr.Button("Extract Highlights", variant="primary") detect_faces_btn = gr.Button("Detect Faces", variant="secondary") status_box = gr.Textbox(label="Status", interactive=False, lines=1) with gr.Column(scale=2): with gr.Row(): clip1 = gr.Video(label="Clip 1", visible=True) clip2 = gr.Video(label="Clip 2", visible=True) clip3 = gr.Video(label="Clip 3", visible=True) with gr.Row(): label1 = gr.Textbox(label="", interactive=False, lines=1) label2 = gr.Textbox(label="", interactive=False, lines=1) label3 = gr.Textbox(label="", interactive=False, lines=1) with gr.Accordion("Transcript", open=False): transcript_box = gr.Textbox( label="Detected lyrics / speech", interactive=False, lines=8, ) with gr.Accordion("Detect Faces (no clips)", open=False): detect_rate = gr.Slider( minimum=0.0, maximum=10.0, value=0.0, step=0.5, label="Sample rate (frames per second, 0 = scan full video)", ) with gr.Row(): check_weights_btn = gr.Button("Check Face Weights", size="sm") extract_face_clips_btn = gr.Button("Extract Face Clips", variant="primary") weights_status = gr.Textbox(label="Face weights status", lines=1, interactive=False) selected_face_status = gr.Textbox( label="Selected face", value="No face selected yet. Click a face thumbnail first.", lines=1, interactive=False, ) clip_duration_face = gr.Number( value=0.0, precision=2, label="Clip duration in seconds (0 = auto face-visible ranges)", ) detections_gallery = gr.Gallery(label="Detected faces", show_label=True) with gr.Row(): face_clip1 = gr.Video(label="Face Clip 1") face_clip2 = gr.Video(label="Face Clip 2") face_clip3 = gr.Video(label="Face Clip 3") face_clip4 = gr.Video(label="Face Clip 4") face_clip5 = gr.Video(label="Face Clip 5") with gr.Row(): face_clip_label1 = gr.Textbox(label="", interactive=False, lines=1) face_clip_label2 = gr.Textbox(label="", interactive=False, lines=1) face_clip_label3 = gr.Textbox(label="", interactive=False, lines=1) face_clip_label4 = gr.Textbox(label="", interactive=False, lines=1) face_clip_label5 = gr.Textbox(label="", interactive=False, lines=1) face_clips_status = gr.Textbox(label="Face clips status", lines=1, interactive=False) detections_json = gr.Textbox(label="Detections JSON", lines=8) def check_face_weights_ui() -> str: st = get_insightface_weight_status() if st["installed"]: return f"Installed: {st['model_dir']}" try: ensure_insightface_weights() st2 = get_insightface_weight_status() if st2["installed"]: return f"Installed after download: {st2['model_dir']}" except Exception as e: return ( f"Install failed: {e} | Missing {len(st['missing'])} file(s) in {st['model_dir']} " f"| Source: {st['source_url']}" ) return f"Still missing files in {st['model_dir']} | Source: {st['source_url']}" def detect_faces_ui(video_path: str, sample_rate: float, progress: gr.Progress = gr.Progress()): no_face_selected_msg = "No face selected yet. Click a face thumbnail first." if not video_path: return ( None, "[]", {}, None, no_face_selected_msg, *_empty_face_clip_outputs("Run face detection first."), ) # Check OpenCV availability early to provide clearer feedback try: import cv2 # noqa: F401 except Exception: return ( None, json.dumps({"error": "OpenCV (cv2) is required for detection. Install with: pip install opencv-python"}), {}, None, no_face_selected_msg, *_empty_face_clip_outputs("Install OpenCV first."), ) try: progress(0.02, desc="Checking InsightFace weights...") # Lazy model bootstrap for environments like Hugging Face Spaces. pre_status = get_insightface_weight_status() if not pre_status["installed"]: progress(0.08, desc="Downloading InsightFace weights...") ensure_insightface_weights() post_status = get_insightface_weight_status() if not post_status["installed"]: return ( None, json.dumps( { "error": "InsightFace weights are still missing after download attempt.", "weights_before": pre_status, "weights_after": post_status, }, indent=2, ), {}, None, no_face_selected_msg, *_empty_face_clip_outputs("Model weights are missing."), ) # Force InsightFace-only detection (no OpenCV fallback). progress(0.20, desc="Loading face detector...") detector = ObjectDetector(use_insightface=True, use_opencv_face_fallback=False) # sample_rate=0 means: inspect all frames in the video effective_rate = float(sample_rate) if effective_rate <= 0: import cv2 cap = cv2.VideoCapture(video_path) fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 cap.release() effective_rate = max(1.0, float(fps)) # Use face detector across the full requested sampling rate identity_similarity_threshold = 0.35 cross_identity_merge_threshold = 0.35 min_face_area = 2500 min_sharpness = 40.0 min_quality_score = 0.08 def on_detect_progress(sampled: int, total_sampled: int, frame_index: int, frame_total: int): denom = max(1, total_sampled) frac = min(1.0, max(0.0, sampled / denom)) progress( 0.20 + 0.78 * frac, desc=f"Detecting faces... {sampled}/{total_sampled} sampled frames", ) res = detector.detect_faces_in_video( video_path, sample_rate=effective_rate, detection_type="face", include_full_frame_fallback=False, progress_callback=on_detect_progress, group_faces=True, identity_similarity_threshold=identity_similarity_threshold, min_face_area=min_face_area, min_sharpness=min_sharpness, min_quality_score=min_quality_score, cross_identity_merge_threshold=cross_identity_merge_threshold, ) except Exception as e: return ( None, json.dumps({"error": str(e)}), {}, None, no_face_selected_msg, *_empty_face_clip_outputs("Face detection failed."), ) frames = res.get("frames", []) identities = res.get("identities", []) identities_json = [] for ident in identities: ident_json = dict(ident) ident_json.pop("embedding_mean", None) identities_json.append(ident_json) # Build gallery images and a stable index->identity map for click selection. images: List[str] = [] gallery_identity_ids: List[Optional[int]] = [] for ident in identities: thumb = ident.get("representative_thumbnail") if thumb: images.append(thumb) gallery_identity_ids.append(int(ident.get("identity_id"))) # Fallback for cases with no grouped identity thumbnails available. if not images: for f in frames: for d in f.get("detections", []): thumb = d.get("thumbnail") if thumb: images.append(thumb) gallery_identity_ids.append(None) payload = { "weights_status": post_status, "identity_similarity_threshold": identity_similarity_threshold, "cross_identity_merge_threshold": cross_identity_merge_threshold, "quality_filter": { "min_face_area": min_face_area, "min_sharpness": min_sharpness, "min_quality_score": min_quality_score, }, "identities": identities_json, "frames": frames, } state_payload = { "video_path": str(video_path), "effective_rate": float(effective_rate), "identities": identities, "gallery_identity_ids": gallery_identity_ids, } progress(1.0, desc="Face detection complete") return ( images or None, json.dumps(payload, indent=2), state_payload, None, no_face_selected_msg, *_empty_face_clip_outputs("Click a face thumbnail, then click Extract Face Clips."), ) def on_face_thumbnail_select(state: Dict[str, Any], evt: gr.SelectData): if not state: return None, "No detection state found. Run Detect Faces first." idx = evt.index if isinstance(idx, (tuple, list)): idx = int(idx[0]) if idx else 0 else: idx = int(idx) identity_map = state.get("gallery_identity_ids", []) if idx < 0 or idx >= len(identity_map): return None, "Invalid selection. Please pick a face thumbnail." identity_id = identity_map[idx] if identity_id is None: return None, "Selected thumbnail is not linked to a grouped identity." identities = state.get("identities", []) selected = None for ident in identities: if int(ident.get("identity_id")) == int(identity_id): selected = ident break if selected is None: return None, "Selected identity not found in detection data." occ = selected.get("occurrences", []) or [] first_ts = ", ".join(format_timestamp(float(o.get("timestamp", 0.0))) for o in occ[:5]) msg = ( f"Selected Face #{int(identity_id)} • detections: {int(selected.get('detections', 0))}" ) if first_ts: msg += f" • first timestamps: {first_ts}" return int(identity_id), msg def extract_face_clips_ui( video_path: str, state: Dict[str, Any], selected_identity_id: Optional[int], clip_duration: Optional[float], progress: gr.Progress = gr.Progress(), ): if not state: return _empty_face_clip_outputs("Run Detect Faces first.") if selected_identity_id is None: return _empty_face_clip_outputs("Click a face thumbnail first.") source_video = str(state.get("video_path") or video_path or "") if not source_video: return _empty_face_clip_outputs("Source video not found.") identities = state.get("identities", []) or [] matched_identities = _collect_matched_identities( identities=identities, selected_identity_id=int(selected_identity_id), similarity_threshold=0.35, ) if not matched_identities: return _empty_face_clip_outputs("Selected face identity is unavailable.") # Collect all occurrences from selected identity + embedding-similar identities. occurrences: List[Dict[str, Any]] = [] matched_ids: List[int] = [] for ident in matched_identities: try: matched_ids.append(int(ident.get("identity_id"))) except Exception: continue occurrences.extend(ident.get("occurrences", []) or []) if not occurrences: return _empty_face_clip_outputs("No timestamp occurrences found for selected face.") duration = float(clip_duration or 0.0) effective_rate = float(state.get("effective_rate") or 1.0) video_duration = _get_video_duration(source_video) segments = _build_face_segments( occurrences=occurrences, clip_duration=duration, video_duration=video_duration, effective_rate=effective_rate, ) if not segments: return _empty_face_clip_outputs("Could not build clip segments from selected face.") out_dir = Path(tempfile.mkdtemp(prefix=f"shortsmith_faceclips_id{int(selected_identity_id):03d}_")) processor = VideoProcessor() videos = [None, None, None, None, None] labels = ["", "", "", "", ""] created = 0 total = len(segments) for i, (start, end) in enumerate(segments): progress( 0.1 + 0.85 * ((i + 1) / max(1, total)), desc=f"Extracting face clip {i+1}/{total}", ) out_path = out_dir / f"face_{int(selected_identity_id):03d}_{i+1:02d}.mp4" try: processor.cut_clip( video_path=source_video, output_path=out_path, start_time=float(start), end_time=float(end), reencode=True, ) created += 1 # Show the first five clips in the UI. All clips are saved to out_dir. if i < 5: videos[i] = str(out_path) labels[i] = f"Face #{int(selected_identity_id)} • {format_timestamp(start)} - {format_timestamp(end)}" except Exception as e: logger.warning(f"Failed to extract face clip {i+1}: {e}") if i < 5: labels[i] = f"Face clip {i+1} failed" status = ( f"Extracted {created}/{total} face clip(s) for Face #{int(selected_identity_id)} " f"(matched IDs: {matched_ids}) • Output dir: {out_dir}" ) progress(1.0, desc="Face clip extraction complete") return ( videos[0], videos[1], videos[2], videos[3], videos[4], labels[0], labels[1], labels[2], labels[3], labels[4], status, ) detect_faces_btn.click( fn=detect_faces_ui, inputs=[video_input, detect_rate], outputs=[ detections_gallery, detections_json, detection_state, selected_identity_state, selected_face_status, face_clip1, face_clip2, face_clip3, face_clip4, face_clip5, face_clip_label1, face_clip_label2, face_clip_label3, face_clip_label4, face_clip_label5, face_clips_status, ], ) detections_gallery.select( fn=on_face_thumbnail_select, inputs=[detection_state], outputs=[selected_identity_state, selected_face_status], ) extract_face_clips_btn.click( fn=extract_face_clips_ui, inputs=[video_input, detection_state, selected_identity_state, clip_duration_face], outputs=[ face_clip1, face_clip2, face_clip3, face_clip4, face_clip5, face_clip_label1, face_clip_label2, face_clip_label3, face_clip_label4, face_clip_label5, face_clips_status, ], ) check_weights_btn.click( fn=check_face_weights_ui, inputs=[], outputs=[weights_status], ) run_btn.click( fn=process_video, inputs=[video_input, num_clips_slider], outputs=[clip1, clip2, clip3, label1, label2, label3, transcript_box, status_box], ) gr.Markdown( """ --- *ShortSmith v4 — runs fully locally, no API keys required.* """ ) if __name__ == "__main__": demo.launch()