import json import os import shutil import time import urllib.request from contextlib import contextmanager from pathlib import Path import cv2 import imageio import matplotlib.pyplot as plt import numpy as np from badas_detector import get_badas_model, run_badas_detector from cosmos_risk_narrator import DEFAULT_REASON_MODEL_NAME, get_reason_model_bundle, run_risk_narrator from extract_clip import extract_pre_alert_clip from predict_backend import get_predict_inference, run_predict_bundle SPACE_ROOT = Path(__file__).resolve().parent CACHE_ROOT = Path(os.environ.get("COSMOS_SPACE_CACHE_DIR") or ("/data/cosmos_sentinel" if Path("/data").exists() else (SPACE_ROOT / ".cache" / "cosmos_sentinel"))) HF_HOME_PATH = Path(os.environ.get("HF_HOME") or ("/data/.huggingface" if Path("/data").exists() else CACHE_ROOT / ".huggingface")) SAMPLE_VIDEO_URL = os.environ.get( "COSMOS_SAMPLE_VIDEO_URL", "https://raw.githubusercontent.com/Ryukijano/Nvidia-Cosmos-Cookoff/main/1_first.mp4", ) PREDICT_OUTPUT_ROOT = CACHE_ROOT / "predict_outputs" PREDICT_MODEL_NAME = os.environ.get("COSMOS_PREDICT_MODEL", "2B/post-trained") os.environ.setdefault("HF_HOME", str(HF_HOME_PATH)) CACHE_ROOT.mkdir(parents=True, exist_ok=True) HF_HOME_PATH.mkdir(parents=True, exist_ok=True) PREDICT_OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) def existing_file(path): if not path: return None resolved = Path(path).resolve() return str(resolved) if resolved.exists() else None @contextmanager def working_directory(path): path = Path(path) path.mkdir(parents=True, exist_ok=True) previous = Path.cwd() os.chdir(path) try: yield path finally: os.chdir(previous) def make_run_dir(prefix="pipeline"): run_dir = CACHE_ROOT / "runs" / f"{prefix}_{time.strftime('%Y%m%d_%H%M%S')}_{int(time.time() * 1000) % 100000}" run_dir.mkdir(parents=True, exist_ok=True) return run_dir def ensure_sample_video(): sample_dir = CACHE_ROOT / "sample_videos" sample_dir.mkdir(parents=True, exist_ok=True) sample_path = sample_dir / "1_first.mp4" if not sample_path.exists(): urllib.request.urlretrieve(SAMPLE_VIDEO_URL, sample_path) return str(sample_path) def cache_uploaded_video(source_path): source_candidate = getattr(source_path, "name", source_path) source = Path(str(source_candidate)) if not source.exists(): raise FileNotFoundError(f"Input video not found: {source}") upload_dir = CACHE_ROOT / "uploads" upload_dir.mkdir(parents=True, exist_ok=True) target = upload_dir / f"{int(time.time())}_{source.name}" shutil.copy2(source, target) return str(target) def preload_runtime(preload_badas=True, preload_reason=True, preload_predict=False, reason_model_name=DEFAULT_REASON_MODEL_NAME, predict_model_name=PREDICT_MODEL_NAME): steps = [] ensure_sample_video() steps.append("Sample video cached") if preload_badas: get_badas_model() steps.append("BADAS model ready") if preload_reason: get_reason_model_bundle(reason_model_name) steps.append(f"Reason model ready: {reason_model_name}") if preload_predict: try: get_predict_inference(predict_model_name, str(PREDICT_OUTPUT_ROOT), True) steps.append(f"Predict model ready: {predict_model_name}") except Exception as e: steps.append(f"Predict model skipped: {e}") return "\n".join(steps) def select_reason_focus_time(badas_result): result = badas_result or {} prediction_window_summary = result.get("prediction_window_summary") or {} top_predictions = result.get("top_predictions") or [] if prediction_window_summary.get("peak_window_end_time") is not None: return float(prediction_window_summary.get("peak_window_end_time")) if top_predictions: return float(top_predictions[0].get("time_sec", 0.0)) return float(result.get("alert_time", 0.0) or 0.0) def extract_frame_at_time(video_path, time_sec): cap = cv2.VideoCapture(video_path) fps = cap.get(cv2.CAP_PROP_FPS) or 0.0 if fps <= 0: cap.release() return None frame_index = max(0, int(round(time_sec * fps))) cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index) ok, frame = cap.read() cap.release() if not ok: return None return frame def apply_full_frame_risk_overlay(frame, intensity, title): if frame is None: return None overlay = frame.copy() heat_color = np.zeros_like(frame) heat_color[:, :] = (0, 0, 255) alpha = max(0.15, min(0.75, float(intensity))) frame = cv2.addWeighted(overlay, 1.0 - alpha, heat_color, alpha, 0) cv2.putText(frame, title, (20, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.85, (255, 255, 255), 2) return frame def build_bbox_heat_overlay(frame, bboxes, title): if frame is None: return None height, width = frame.shape[:2] heat = np.zeros((height, width), dtype=np.float32) rendered = frame.copy() for label, bbox in (bboxes or {}).items(): if not isinstance(bbox, (list, tuple)) or len(bbox) != 4: continue x1, y1, x2, y2 = bbox px1 = max(0, min(width - 1, int(round(x1 * width)))) py1 = max(0, min(height - 1, int(round(y1 * height)))) px2 = max(px1 + 1, min(width, int(round(x2 * width)))) py2 = max(py1 + 1, min(height, int(round(y2 * height)))) heat[py1:py2, px1:px2] += 1.0 cv2.rectangle(rendered, (px1, py1), (px2, py2), (0, 255, 0), 2) cv2.putText(rendered, str(label), (px1, max(20, py1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) if heat.max() > 0: heat = heat / heat.max() heat_u8 = np.uint8(255 * heat) heat_color = cv2.applyColorMap(heat_u8, cv2.COLORMAP_JET) rendered = cv2.addWeighted(rendered, 0.65, heat_color, 0.35, 0) cv2.putText(rendered, title, (20, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.85, (255, 255, 255), 2) return rendered def save_frame_strip(frames, output_path, resize_height=220): valid_frames = [frame for frame in frames if frame is not None] if not valid_frames: return None resized = [] for frame in valid_frames: height, width = frame.shape[:2] scale = resize_height / max(height, 1) resized.append(cv2.resize(frame, (max(1, int(round(width * scale))), resize_height))) strip = cv2.hconcat(resized) cv2.imwrite(str(output_path), strip) return existing_file(output_path) def create_badas_frame_strip(video_path, badas_result, output_path): top_predictions = (badas_result or {}).get("top_predictions") or [] if not top_predictions: return None frames = [] for item in top_predictions[:4]: frame = extract_frame_at_time(video_path, float(item.get("time_sec", 0.0))) if frame is None: continue frames.append( apply_full_frame_risk_overlay( frame, float(item.get("probability", 0.0)), f"BADAS {item.get('time_sec', 0.0):.2f}s | {item.get('probability', 0.0):.1%}", ) ) return save_frame_strip(frames, output_path) def create_reason_frame_strip(clip_path, reason_payload, output_path): frame_metadata = (reason_payload or {}).get("frame_metadata") or {} timestamps = frame_metadata.get("sampled_timestamps_sec") or [] if not timestamps: return None bboxes = (reason_payload or {}).get("bboxes") or {} if not bboxes: return None frames = [] for timestamp in timestamps[:4]: frame = extract_frame_at_time(clip_path, float(timestamp)) if frame is None: continue frames.append(build_bbox_heat_overlay(frame, bboxes, f"Reason {float(timestamp):.2f}s | bbox focus")) return save_frame_strip(frames, output_path) def create_visualizations(source_video_path, clip_path, badas_result, reason_payload): bboxes = (reason_payload or {}).get("bboxes") or {} risk_score = (reason_payload or {}).get("risk_score") or 0 bbox_image = Path("bboxes_visualization.png") risk_image = Path("risk_visualization.png") overlay_gif = Path("video_with_bboxes.gif") badas_strip_image = Path("badas_frame_strip.png") reason_strip_image = Path("reason_frame_strip.png") if bboxes: fig, ax = plt.subplots(figsize=(6, 6)) for label, bbox in bboxes.items(): x1, y1, x2, y2 = bbox rect = plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="red", linewidth=2) ax.add_patch(rect) ax.text(x1, y1, label, fontsize=12, color="red") ax.set_xlim(0, 1) ax.set_ylim(0, 1) ax.set_title("Detected Agents Bounding Boxes") ax.invert_yaxis() plt.savefig(bbox_image) plt.close() fig, ax = plt.subplots(figsize=(4, 2)) ax.barh(["Risk Score"], [risk_score], color="orange") ax.set_xlim(0, 5) ax.set_title("Collision Risk Assessment") plt.savefig(risk_image) plt.close() if bboxes: cap = cv2.VideoCapture(str(clip_path)) frames = [] frame_count = 0 max_frames = 20 while cap.isOpened() and frame_count < max_frames: ret, frame = cap.read() if not ret: break for label, bbox in bboxes.items(): x1, y1, x2, y2 = bbox height, width = frame.shape[:2] cv2.rectangle(frame, (int(x1 * width), int(y1 * height)), (int(x2 * width), int(y2 * height)), (0, 255, 0), 2) cv2.putText(frame, label, (int(x1 * width), int(y1 * height) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) frame_count += 1 cap.release() if frames: imageio.mimsave(overlay_gif, frames, fps=5, loop=0) badas_strip = create_badas_frame_strip(source_video_path, badas_result, badas_strip_image) reason_strip = create_reason_frame_strip(clip_path, reason_payload, reason_strip_image) return { "bbox_image": existing_file(bbox_image) if bboxes else None, "risk_image": existing_file(risk_image), "overlay_gif": existing_file(overlay_gif) if bboxes else None, "badas_frame_strip": badas_strip, "reason_frame_strip": reason_strip, } def build_pipeline_overview(badas_result, reason_payload): threshold_summary = (badas_result or {}).get("threshold_summary") or {} prediction_window_summary = (badas_result or {}).get("prediction_window_summary") or {} parsing_summary = (reason_payload or {}).get("parsing_summary") or {} frame_metadata = (reason_payload or {}).get("frame_metadata") or {} validation = (reason_payload or {}).get("validation") or {} validation_flags = validation.get("flags") or {} return { "collision_gate_triggered": bool((badas_result or {}).get("collision_detected")), "alert_time_sec": (badas_result or {}).get("alert_time"), "reason_focus_time_sec": select_reason_focus_time(badas_result), "alert_confidence": (badas_result or {}).get("confidence"), "threshold_crossing_count": threshold_summary.get("threshold_crossing_count", 0), "peak_probability": (badas_result or {}).get("valid_prediction_max"), "peak_window_average_probability": prediction_window_summary.get("max_average_probability"), "incident_type": (reason_payload or {}).get("incident_type"), "severity_label": (reason_payload or {}).get("severity_label"), "reason_risk_score": (reason_payload or {}).get("risk_score"), "reason_bbox_count": (reason_payload or {}).get("bbox_count", 0), "reason_prompt_conditioned_by_badas": bool((reason_payload or {}).get("badas_context")), "reason_missing_fields": parsing_summary.get("missing_fields", []), "reason_processed_frame_count": frame_metadata.get("processed_frame_count"), "reason_output_reliable": validation.get("is_reliable"), "reason_second_pass_used": validation_flags.get("second_pass_used", False), "reason_fallback_override_applied": validation_flags.get("fallback_override_applied", False), } def build_reason_payload(video_path, focus_video_path, badas_context): result_text, metadata = run_risk_narrator(video_path, badas_context=badas_context, focus_video_path=focus_video_path) payload = metadata.get("parsed_payload") or {} payload["video_path"] = video_path payload["focus_video_path"] = focus_video_path payload["user_prompt"] = metadata["user_prompt"] payload["badas_context"] = metadata["badas_context"] payload["frame_metadata"] = metadata["frame_metadata"] payload["focus_frame_metadata"] = metadata["focus_frame_metadata"] payload["video_input_count"] = metadata["video_input_count"] payload["model_metadata"] = metadata["model"] payload["generation_config"] = metadata["generation_config"] payload["input_token_count"] = metadata["input_token_count"] payload["output_token_count"] = metadata["output_token_count"] payload["text"] = payload.get("text") or result_text return payload def run_pipeline(video_path, include_predict=False, predict_modes=None, predict_model_name=PREDICT_MODEL_NAME): run_dir = make_run_dir("pipeline") log_lines = ["🚀 Starting Cosmos Sentinel Gradio pipeline", f"Input video: {video_path}"] with working_directory(run_dir): log_lines.append("📍 Step 1: BADAS V-JEPA2 Collision Detection") badas_result = run_badas_detector(video_path) log_lines.append("📍 Step 2: Extracting Pre-Alert Clip") reason_focus_time = select_reason_focus_time(badas_result) extracted_clip = extract_pre_alert_clip(video_path, reason_focus_time, "./extracted_clip.mp4") if not extracted_clip: raise RuntimeError("Failed to extract BADAS-focused clip") log_lines.append("📍 Step 3: Cosmos Reason 2 Risk Analysis") reason_payload = build_reason_payload(video_path, extracted_clip, badas_result) visualizations = create_visualizations(video_path, extracted_clip, badas_result, reason_payload) predict_payload = None if include_predict: log_lines.append("📍 Step 4: Cosmos Predict continuation") selected_modes = predict_modes or ["prevented_continuation", "observed_continuation"] predict_payload = run_predict_bundle( video_path, badas_context=badas_result, reason_context=reason_payload, modes=selected_modes, model_name=predict_model_name, output_root=PREDICT_OUTPUT_ROOT / run_dir.name, fallback_conditioning_path=extracted_clip, ) pipeline_payload = { "input_video": video_path, "pipeline_mode": "badas_reason_predict" if include_predict else "badas_reason_only", "iterations": [ { "iteration": 1, "input_video": video_path, "steps": { "badas": { "success": True, "alert_time": badas_result.get("alert_time"), "reason_focus_time": reason_focus_time, "result": badas_result, }, "clip_extraction": { "success": True, "clip_path": existing_file(extracted_clip), "alert_time": badas_result.get("alert_time"), "reason_focus_time": reason_focus_time, }, "reason": { "success": True, "full_video_input": video_path, "focus_clip_input": existing_file(extracted_clip), "result": reason_payload, "text": reason_payload.get("text", ""), "visualizations": visualizations, }, }, } ], "artifacts": { "extracted_clip": existing_file(extracted_clip), "badas_gradient_saliency": existing_file((badas_result or {}).get("gradient_saliency_image")), "bbox_image": visualizations.get("bbox_image"), "risk_image": visualizations.get("risk_image"), "overlay_gif": visualizations.get("overlay_gif"), "badas_frame_strip": visualizations.get("badas_frame_strip"), "reason_frame_strip": visualizations.get("reason_frame_strip"), }, "status": "completed", "overview": build_pipeline_overview(badas_result, reason_payload), "run_directory": str(run_dir), } if predict_payload: pipeline_payload["predict"] = predict_payload for artifact_key, artifact_value in (predict_payload.get("artifacts") or {}).items(): pipeline_payload["artifacts"][artifact_key] = artifact_value log_lines.append("🎉 Cosmos Sentinel pipeline completed") return { "success": True, "logs": "\n".join(log_lines), "pipeline_payload": pipeline_payload, "badas_result": badas_result, "reason_result": reason_payload, "predict_payload": predict_payload, "run_directory": str(run_dir), } def run_predict_only(pipeline_payload, selection="both", predict_model_name=PREDICT_MODEL_NAME): if not pipeline_payload: raise ValueError("Run BADAS + Reason before Predict") iteration = ((pipeline_payload.get("iterations") or [{}])[-1]) steps = iteration.get("steps") or {} badas_result = (steps.get("badas") or {}).get("result") or {} reason_result = (steps.get("reason") or {}).get("result") or {} artifacts = (pipeline_payload.get("artifacts") or {}) source_video = pipeline_payload.get("input_video") modes = ["prevented_continuation", "observed_continuation"] if selection == "both" else [selection] run_dir = make_run_dir("predict") predict_payload = run_predict_bundle( source_video, badas_context=badas_result, reason_context=reason_result, modes=modes, model_name=predict_model_name, output_root=PREDICT_OUTPUT_ROOT / run_dir.name, fallback_conditioning_path=artifacts.get("extracted_clip"), ) merged = json.loads(json.dumps(pipeline_payload)) merged["predict"] = predict_payload merged_artifacts = merged.get("artifacts") or {} for artifact_key, artifact_value in (predict_payload.get("artifacts") or {}).items(): merged_artifacts[artifact_key] = artifact_value merged["artifacts"] = merged_artifacts return predict_payload, merged