import torch from transformers import Qwen3VLForConditionalGeneration, AutoProcessor import cv2 from PIL import Image import ast import json import os import re import sys from functools import lru_cache import numpy as np DEFAULT_REASON_MODEL_NAME = os.environ.get("COSMOS_REASON_MODEL", "nvidia/Cosmos-Reason2-8B") @lru_cache(maxsize=2) def get_reason_model_bundle(model_name=DEFAULT_REASON_MODEL_NAME): print(f"Loading base Cosmos Reason model: {model_name}") model = Qwen3VLForConditionalGeneration.from_pretrained( model_name, device_map="auto", torch_dtype=torch.float16, ) print("Using pretrained checkpoint only; adapter loading is disabled") processor = AutoProcessor.from_pretrained(model_name) print("✅ Fine-tuned model loaded (FP16) on", model.device) return model_name, model, processor # System and user prompts for CCTV traffic reasoning SYSTEM_PROMPT = "You are Cosmos Risk Narrator, an AI traffic-safety analyst reading fixed CCTV or roadside camera footage. Your task is to detect possible accidents, distinguish near-misses from confirmed impacts, and estimate crash severity in clear operational language." USER_PROMPT_TEMPLATE = """Analyze the CCTV traffic video for collision and severity assessment. Focus on: - Traffic environment analysis including road geometry, density, visibility, and possible occlusions from the CCTV view - Distinguishing no incident, near-miss, minor collision, serious collision, or multi-agent severe crash - Severity scoring (0-5 scale) where 0=no incident, 1=low concern, 2=near miss, 3=minor collision risk/contact, 4=serious collision, 5=catastrophic or multi-vehicle severe crash - Short causal explanation describing what makes the event severe or not severe - Treat visible physical vehicle-to-vehicle contact, deformation, abrupt rebound, spin, or forced redirection after contact as evidence of a real collision rather than a near miss - Use near_miss only if the agents clearly avoid contact in the visible frames - If two cars visibly crash into each other in the clip, do not output no_incident - Focus especially on frames near the BADAS alert time and the seconds immediately after it Use the BADAS detector summary below as a high-priority cue for where risk may be developing, but verify it against the actual video frames before concluding that a crash happened. Video inputs: - Video 1 is the full CCTV video and should be treated as the primary source of truth for temporal context - Video 2, if present, is a BADAS-focused evidence clip centered on the highest-risk moment and should be used to inspect the likely impact sequence more closely Important reliability rule: precise visual grounding bounding boxes are not reliable from this video prompt alone. Do not invent placeholder boxes. If you cannot localize agents reliably from the sampled video frames, return an empty bbox set. BADAS detector context: {badas_context} Output format: Scene summary: [1-2 sentence summary of what the CCTV camera sees] Incident type: [no_incident | near_miss | collision | multi_vehicle_collision | unclear] Severity label: [none | low | moderate | high | critical] Bounding boxes: [brief text summary only, or none] Critical risk detected at: X.X seconds Risk score: X/5 At-risk agent: [description] Explanation: [clear sentence about why this is dangerous] Time-to-impact if no action: X.X seconds Bounding boxes (normalized 0-1): - Return no lines here unless localization is genuinely reliable Counterfactual collision prompt: [1-sentence prompt describing the collision that would occur if the driver didn't brake]""" USER_PROMPT_TEMPLATE = USER_PROMPT_TEMPLATE + "\n\nReturn exactly one report block in the format above. Do not repeat the report." SECOND_PASS_PROMPT_TEMPLATE = """Re-evaluate this CCTV clip only for collision confirmation and severity. Rules: - Decide whether visible contact or immediate post-contact forced redirection occurs. - If contact, rebound, spin, impact overlap, or abrupt trajectory change after contact is visible, output collision or multi_vehicle_collision. - If the clip is ambiguous, output unclear. - Do not output no_incident if BADAS flagged a strong collision cue and the clip shows a plausible impact sequence. - Do not output any bounding boxes. Leave the normalized bounding box section empty. Video inputs: - Video 1 is the full CCTV video - Video 2, if present, is the BADAS-focused evidence clip around the likely impact window BADAS detector context: {badas_context} Output format: Scene summary: [1-2 sentence summary of what the CCTV camera sees] Incident type: [no_incident | near_miss | collision | multi_vehicle_collision | unclear] Severity label: [none | low | moderate | high | critical] Bounding boxes: none Critical risk detected at: X.X seconds Risk score: X/5 At-risk agent: [description] Explanation: [clear sentence about why this is dangerous] Time-to-impact if no action: X.X seconds or N/A Bounding boxes (normalized 0-1): Counterfactual collision prompt: [1 sentence]""" SECOND_PASS_PROMPT_TEMPLATE = SECOND_PASS_PROMPT_TEMPLATE + "\n\nReturn exactly one report block in the format above. Do not repeat the report." GENERATION_CONFIG = { "max_new_tokens": 1024, "do_sample": False, } VIDEO_BBOX_SUPPORT_NOTE = "Qwen3-VL video prompting does not reliably support native per-video bbox grounding in this pipeline, so bbox output is treated as untrusted unless sourced from a dedicated image-grounding path." def extract_sections(text): sections = { "scene_summary_line": "", "incident_type_line": "", "severity_label_line": "", "bounding_boxes_header": "", "critical_risk_line": "", "risk_score_line": "", "at_risk_agent_line": "", "explanation_line": "", "time_to_impact_line": "", "counterfactual_line": "", } for line in text.splitlines(): stripped = line.strip() if stripped.startswith("Scene summary:"): sections["scene_summary_line"] = stripped elif stripped.startswith("Incident type:"): sections["incident_type_line"] = stripped elif stripped.startswith("Severity label:"): sections["severity_label_line"] = stripped elif stripped.startswith("Bounding boxes:"): sections["bounding_boxes_header"] = stripped elif stripped.startswith("Critical risk detected at:"): sections["critical_risk_line"] = stripped elif stripped.startswith("Risk score:"): sections["risk_score_line"] = stripped elif stripped.startswith("At-risk agent:"): sections["at_risk_agent_line"] = stripped elif stripped.startswith("Explanation:"): sections["explanation_line"] = stripped elif stripped.startswith("Time-to-impact if no action:"): sections["time_to_impact_line"] = stripped elif stripped.startswith("Counterfactual collision prompt:"): sections["counterfactual_line"] = stripped return sections def extract_reason_payload(text): payload = { "text": text, "system_prompt": SYSTEM_PROMPT, "user_prompt": "", "badas_context": {}, "scene_summary": "", "incident_type": "", "severity_label": "", "risk_score": None, "critical_risk_time": None, "at_risk_agent": "", "explanation": "", "time_to_impact": None, "counterfactual_prompt": "", "bboxes": {}, } sections = extract_sections(text) bbox_section = False for line in text.splitlines(): stripped = line.strip() if not stripped: if bbox_section: bbox_section = False continue if stripped == "Bounding boxes (normalized 0-1):": bbox_section = True continue if bbox_section and ":" in stripped: label, coords_text = stripped.split(":", 1) try: coords = ast.literal_eval(coords_text.strip()) if isinstance(coords, (list, tuple)) and len(coords) == 4: payload["bboxes"][label.strip()] = [float(value) for value in coords] except (ValueError, SyntaxError): pass if stripped.startswith("Scene summary:"): payload["scene_summary"] = stripped.split(":", 1)[1].strip() elif stripped.startswith("Incident type:"): payload["incident_type"] = stripped.split(":", 1)[1].strip() elif stripped.startswith("Severity label:"): payload["severity_label"] = stripped.split(":", 1)[1].strip() elif stripped.startswith("Risk score:"): match = re.search(r"(\d+)", stripped) if match: payload["risk_score"] = int(match.group(1)) elif stripped.startswith("Critical risk detected at:"): match = re.search(r"(\d+(?:\.\d+)?)", stripped) if match: payload["critical_risk_time"] = float(match.group(1)) elif stripped.startswith("At-risk agent:"): payload["at_risk_agent"] = stripped.split(":", 1)[1].strip() elif stripped.startswith("Explanation:"): payload["explanation"] = stripped.split(":", 1)[1].strip() elif stripped.startswith("Time-to-impact if no action:"): match = re.search(r"(\d+(?:\.\d+)?)", stripped) if match: payload["time_to_impact"] = float(match.group(1)) elif stripped.startswith("Counterfactual collision prompt:"): payload["counterfactual_prompt"] = stripped.split(":", 1)[1].strip() parsed_fields = { "scene_summary": bool(payload["scene_summary"]), "incident_type": bool(payload["incident_type"]), "severity_label": bool(payload["severity_label"]), "risk_score": payload["risk_score"] is not None, "critical_risk_time": payload["critical_risk_time"] is not None, "at_risk_agent": bool(payload["at_risk_agent"]), "explanation": bool(payload["explanation"]), "time_to_impact": payload["time_to_impact"] is not None, "counterfactual_prompt": bool(payload["counterfactual_prompt"]), "bboxes": bool(payload["bboxes"]), } payload["bbox_count"] = len(payload["bboxes"]) payload["bbox_labels"] = list(payload["bboxes"].keys()) payload["sections"] = sections payload["parsing_summary"] = { "parsed_field_flags": parsed_fields, "missing_fields": [field for field, parsed in parsed_fields.items() if not parsed], "parsed_field_count": int(sum(1 for parsed in parsed_fields.values() if parsed)), "total_expected_fields": len(parsed_fields), } return payload def payload_requires_second_pass(payload, badas_context): if not payload: return True incident_type = (payload.get("incident_type") or "").strip().lower() risk_score = payload.get("risk_score") badas_confidence = float((badas_context or {}).get("confidence") or 0.0) peak_probability = float((badas_context or {}).get("valid_prediction_max") or 0.0) collision_detected = bool((badas_context or {}).get("collision_detected")) explanation = (payload.get("explanation") or "").strip().lower() missing_core = not payload.get("scene_summary") or not payload.get("incident_type") or not payload.get("severity_label") undercalls_badas = incident_type == "no_incident" and (collision_detected or badas_confidence >= 0.45 or peak_probability >= 0.6) noncommittal_zero_risk = incident_type in {"", "no_incident"} and (risk_score in {None, 0}) and "no immediate risk" in explanation return bool(missing_core or undercalls_badas or noncommittal_zero_risk) def apply_badas_consistency_fallback(payload, badas_context): payload = dict(payload or {}) incident_type = (payload.get("incident_type") or "").strip().lower() confidence = float((badas_context or {}).get("confidence") or 0.0) peak_probability = float((badas_context or {}).get("valid_prediction_max") or 0.0) peak_window_average = float(((badas_context or {}).get("prediction_window_summary") or {}).get("max_average_probability") or 0.0) threshold_crossings = int((((badas_context or {}).get("threshold_summary") or {}).get("threshold_crossing_count") or 0)) strong_badas_evidence = bool((badas_context or {}).get("collision_detected")) and ( confidence >= 0.5 or peak_probability >= 0.7 or peak_window_average >= 0.6 or threshold_crossings >= 2 ) if incident_type != "no_incident" or not strong_badas_evidence: payload.setdefault("fallback_override", None) return payload focus_time = float((badas_context or {}).get("alert_time") or 0.0) payload["incident_type"] = "collision" payload["severity_label"] = payload.get("severity_label") if payload.get("severity_label") in {"high", "critical"} else "moderate" payload["risk_score"] = max(int(payload.get("risk_score") or 0), 3) payload["at_risk_agent"] = payload.get("at_risk_agent") or "vehicles in the BADAS high-risk collision window" payload["critical_risk_time"] = payload.get("critical_risk_time") if payload.get("critical_risk_time") is not None else focus_time explanation = (payload.get("explanation") or "").strip() if not explanation or "no immediate risk" in explanation.lower(): payload["explanation"] = ( f"BADAS detected a high-confidence collision sequence near {focus_time:.2f}s, so the inconsistent no_incident response was escalated to collision for safety-focused review." ) payload["fallback_override"] = { "applied": True, "source": "badas_consistency_guard", "reason": "Reason output conflicted with strong BADAS collision evidence", } return payload def attach_validation(payload, badas_context, second_pass_used=False, initial_payload=None): payload = apply_badas_consistency_fallback(payload, badas_context) payload = dict(payload or {}) validation_flags = { "video_bbox_grounding_supported": False, "reason_bboxes_rejected": bool(payload.get("bboxes")), "second_pass_used": bool(second_pass_used), "initial_payload_replaced": bool(second_pass_used), "incident_conflicts_with_badas": False, "fallback_override_applied": bool((payload.get("fallback_override") or {}).get("applied")), } incident_type = (payload.get("incident_type") or "").strip().lower() badas_confidence = float((badas_context or {}).get("confidence") or 0.0) peak_probability = float((badas_context or {}).get("valid_prediction_max") or 0.0) if incident_type == "no_incident" and (bool((badas_context or {}).get("collision_detected")) or badas_confidence >= 0.45 or peak_probability >= 0.6): validation_flags["incident_conflicts_with_badas"] = True payload["raw_bboxes"] = payload.get("bboxes") or {} payload["bboxes"] = {} payload["bbox_count"] = 0 payload["bbox_labels"] = [] payload["validation"] = { "is_reliable": not validation_flags["incident_conflicts_with_badas"], "flags": validation_flags, "note": VIDEO_BBOX_SUPPORT_NOTE, } return payload def summarize_badas_context(badas_context): if not badas_context: return "No BADAS detector context was provided. Infer risk only from the visible clip." threshold_summary = badas_context.get("threshold_summary") or {} prediction_window_summary = badas_context.get("prediction_window_summary") or {} top_predictions = badas_context.get("top_predictions") or [] top_prediction_lines = [] for item in top_predictions[:3]: top_prediction_lines.append( f"- sampled_frame={item.get('sampled_frame')} time={item.get('time_sec', 0.0):.2f}s probability={item.get('probability', 0.0):.2%}" ) if not top_prediction_lines: top_prediction_lines.append("- no top prediction frames available") return "\n".join([ f"collision_detected={bool(badas_context.get('collision_detected'))}", f"alert_time={float(badas_context.get('alert_time', 0.0)):.2f}s", f"alert_confidence={float(badas_context.get('confidence', 0.0)):.2%}", f"threshold={float(badas_context.get('threshold', 0.0)):.2f}", f"threshold_crossing_count={int(threshold_summary.get('threshold_crossing_count', 0))}", f"peak_probability={float(badas_context.get('valid_prediction_max', 0.0) or 0.0):.2%}", f"peak_window_average_probability={float(prediction_window_summary.get('max_average_probability', 0.0) or 0.0):.2%}", "top_prediction_frames:", *top_prediction_lines, ]) def build_user_prompt(badas_context): return USER_PROMPT_TEMPLATE.format(badas_context=summarize_badas_context(badas_context)) def build_second_pass_prompt(badas_context): return SECOND_PASS_PROMPT_TEMPLATE.format(badas_context=summarize_badas_context(badas_context)) def process_video(video_path, max_frames=96): cap = cv2.VideoCapture(video_path) fps = cap.get(cv2.CAP_PROP_FPS) target_fps = 4.0 frame_interval = max(1, int(round(fps / target_fps))) if fps else 1 total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) frames = [] count = 0 sampled_indices_all = [] while True: ret, frame = cap.read() if not ret: break if count % frame_interval == 0: rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(Image.fromarray(rgb)) sampled_indices_all.append(int(count)) count += 1 cap.release() if len(frames) > max_frames: selected_positions = np.linspace(0, len(frames) - 1, max_frames, dtype=int).tolist() frames = [frames[position] for position in selected_positions] sampled_indices = [sampled_indices_all[position] for position in selected_positions] else: sampled_indices = sampled_indices_all print(f"Processed {len(frames)} frames") return frames, { "video_path": video_path, "original_fps": float(fps) if fps else 0.0, "target_fps": float(target_fps), "frame_interval": int(frame_interval), "total_frames": int(total_frames), "duration_sec": float(total_frames / fps) if fps else 0.0, "processed_frame_count": int(len(frames)), "sampled_frame_indices": sampled_indices, "sampled_timestamps_sec": [float(index / fps) for index in sampled_indices] if fps else [], "max_frames": int(max_frames), } def generate_reason_response(video_inputs, user_prompt, model_name=None): resolved_model_name, model, processor = get_reason_model_bundle(model_name or DEFAULT_REASON_MODEL_NAME) content = [] for video_input in video_inputs: content.append( { "type": "video", "video": video_input["frames"], "fps": float(video_input["fps"]), "max_pixels": 768 * 768, } ) content.append({"type": "text", "text": user_prompt}) messages = [ {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}, {"role": "user", "content": content} ] inputs = processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt", add_vision_id=True, ).to(model.device) print("Generating risk assessment...") generated = model.generate( **inputs, **GENERATION_CONFIG, ) output = processor.batch_decode(generated, skip_special_tokens=True)[0] assistant_part = output.split("assistant")[-1].strip() return assistant_part, inputs, generated, { "base_model_name": resolved_model_name, "device": str(model.device), "dtype": str(getattr(model, "dtype", torch.float16)), } def run_risk_narrator(video_path, badas_context=None, focus_video_path=None, model_name=None): frames, frame_metadata = process_video(video_path) user_prompt = build_user_prompt(badas_context or {}) video_inputs = [ { "label": "full_video", "frames": frames, "fps": float(frame_metadata.get("target_fps") or 4.0), "metadata": frame_metadata, } ] focus_frame_metadata = None if focus_video_path: focus_frames, focus_frame_metadata = process_video(focus_video_path, max_frames=48) if focus_frames: video_inputs.append( { "label": "focus_clip", "frames": focus_frames, "fps": float(focus_frame_metadata.get("target_fps") or 4.0), "metadata": focus_frame_metadata, } ) assistant_part, inputs, generated, model_metadata = generate_reason_response(video_inputs, user_prompt, model_name=model_name) initial_payload = extract_reason_payload(assistant_part) second_pass_used = False if payload_requires_second_pass(initial_payload, badas_context or {}): second_pass_used = True second_pass_prompt = build_second_pass_prompt(badas_context or {}) assistant_part, inputs, generated, model_metadata = generate_reason_response(video_inputs, second_pass_prompt, model_name=model_name) user_prompt = second_pass_prompt final_payload = extract_reason_payload(assistant_part) else: final_payload = initial_payload final_payload = attach_validation(final_payload, badas_context or {}, second_pass_used=second_pass_used, initial_payload=initial_payload) return assistant_part, { "parsed_payload": final_payload, "frame_metadata": frame_metadata, "focus_frame_metadata": focus_frame_metadata, "video_input_count": len(video_inputs), "user_prompt": user_prompt, "badas_context": badas_context or {}, "model": { "base_model_name": model_metadata["base_model_name"], "checkpoint_source": "pretrained_huggingface_checkpoint", "adapter_loaded": False, "device": model_metadata["device"], "dtype": model_metadata["dtype"], }, "generation_config": GENERATION_CONFIG, "input_token_count": int(inputs["input_ids"].shape[-1]), "output_token_count": int(generated.shape[-1]), } if __name__ == "__main__": video_path = sys.argv[1] if len(sys.argv) > 1 else "./extracted_clip.mp4" focus_video_path = sys.argv[2] if len(sys.argv) > 2 else None badas_context_raw = os.environ.get("COSMOS_BADAS_CONTEXT", "") try: badas_context = json.loads(badas_context_raw) if badas_context_raw else {} except json.JSONDecodeError: badas_context = {} result, metadata = run_risk_narrator(video_path, badas_context=badas_context, focus_video_path=focus_video_path) payload = metadata.get("parsed_payload") or extract_reason_payload(result) 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"] print("Reason 2 Output:") print(result) if payload["counterfactual_prompt"]: counterfactual = payload["counterfactual_prompt"] print(f"Counterfactual Prompt: {counterfactual}") else: print("Counterfactual prompt not found") print(f"REASON_JSON: {json.dumps(payload)}")