""" AiGiS Safety Captioning — Hugging Face Space (Gradio UI). Three modes, chosen from the Mode dropdown: * Full pipeline : YOLO checks PPE per person (deterministic) and the VLM writes the safety record: issue, main hazard, missing controls and a safety observation. The detector stays authoritative for PPE; the VLM reports its findings in plain words and flags any the image clearly contradicts. * Detector only : run YOLO, show boxes + raw detections, no VLM (no API cost). * VLM only : skip YOLO entirely; the VLM judges the image on its own. The user also picks which VLM (Claude or Qwen3-VL) writes the hazard statement. The YOLO weights are pulled once at startup from a Hugging Face model repo. Configure via Space variables: - AIGIS_HF_REPO_ID (e.g. ZahraFatima/aigis-yolo) - AIGIS_HF_FILENAME (e.g. best.pt) - AIGIS_CONF detection confidence threshold (default 0.15) - AIGIS_TTA "1" to enable test-time augmentation (slower, higher recall) Which PPE is required or advisory, and the confidence a box needs to count as compliance evidence, are set in sh17.yaml (ppe_rules, compliance_conf). Secrets required: - ANTHROPIC_API_KEY for Claude models - DASHSCOPE_API_KEY for Qwen models """ from __future__ import annotations import os import cv2 import gradio as gr import numpy as np import spaces from huggingface_hub import hf_hub_download import vlm_caption from detector import SafetyDetector HF_REPO_ID = os.environ.get("AIGIS_HF_REPO_ID", "ZahraFatima/aigis-yolo") HF_FILENAME = os.environ.get("AIGIS_HF_FILENAME", "best.pt") DATA_YAML = os.environ.get("AIGIS_DATA_YAML", "sh17.yaml") CONF = float(os.environ.get("AIGIS_CONF", "0.15")) TTA = os.environ.get("AIGIS_TTA", "0") == "1" MODE_FULL = "Full pipeline (YOLO + VLM)" MODE_DETECTOR = "Detector only (no VLM, no API cost)" MODE_VLM = "VLM only (no YOLO)" _PROVIDER_KEY = {"anthropic": "ANTHROPIC_API_KEY", "dashscope": "DASHSCOPE_API_KEY"} weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME) detector = SafetyDetector(weights_path, DATA_YAML, conf=CONF, augment=TTA) def _key_error_for(model: str, annotated, det_dict): _, provider, _ = vlm_caption._resolve(model) needed_key = _PROVIDER_KEY.get(provider) if needed_key and needed_key not in os.environ: return ( annotated, f"**Error:** {needed_key} is not set for this Space " f"(required for model '{model}').", det_dict, ) return None @spaces.GPU(duration=60) def run(image: np.ndarray, mode: str, model: str): if image is None: return None, "Upload an image first.", {} # ---- VLM only: skip YOLO entirely ---- if mode == MODE_VLM: err = _key_error_for(model, image, {}) if err: return err frame_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) ok, buf = cv2.imencode(".jpg", frame_bgr) vlm_result = vlm_caption.caption_bytes( buf.tobytes(), "image/jpeg", detection_result=None, model=model, use_detector=False, ) caption_md = _format_caption(vlm_result, model, image_only=True) return image, caption_md, {"note": "VLM-only mode: no detector was run."} # ---- Modes that use YOLO ---- frame_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) det_result, annotated_bgr = detector.detect_and_annotate(frame_bgr, name="upload.jpg") annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB) det_dict = det_result.to_dict() if mode == MODE_DETECTOR: return annotated_rgb, "_(Detector only — VLM captioning skipped)_", det_dict # ---- Full pipeline ---- err = _key_error_for(model, annotated_rgb, det_dict) if err: return err ok, buf = cv2.imencode(".jpg", frame_bgr) vlm_result = vlm_caption.caption_bytes( buf.tobytes(), "image/jpeg", detection_result=det_dict, model=model, use_detector=True, ) caption_md = _format_caption(vlm_result, model, image_only=False) return annotated_rgb, caption_md, det_dict def _format_caption(vlm_result: dict, model: str, image_only: bool) -> str: header = f"**Model:** {vlm_result.get('_model', model)}" if image_only: header += " · _image-only (no detector)_" if vlm_result.get("parse_error"): return f"{header}\n\n**Error:** model did not return valid JSON.\n\n{vlm_result.get('raw', '')}" md = header # Hazard first, in the client's order: issue, main hazard, missing controls. issue = vlm_result.get("issue", "") if issue: md += f"\n\n**Issue:** {issue}" main_hazard = vlm_result.get("main_hazard", "") if main_hazard: md += f"\n\n**Main hazard:** {main_hazard}" controls = vlm_result.get("missing_controls") or [] if isinstance(controls, str): controls = [controls] if controls: md += "\n\n**Missing controls:** " + "; ".join(str(c) for c in controls) md += ( f"\n\n**Safety observation:** {vlm_result.get('safety_observation', '')}\n\n" f"**Severity:** {vlm_result.get('severity', 'n/a')}" ) disagree = vlm_result.get("detector_disagreement", "") if disagree: md += f"\n\n**⚠ Detector/image disagreement:** {disagree}" # Fixed reminder, so the model does not spend a sentence on it every time. if image_only: md += "\n\n_Assessed from the image alone. Verify on site._" else: md += "\n\n_Assessed from a single image. Verify on site._" return md demo = gr.Interface( fn=run, inputs=[ gr.Image(type="numpy", label="Upload image"), gr.Dropdown( choices=[MODE_FULL, MODE_DETECTOR, MODE_VLM], value=MODE_FULL, label="Mode", ), gr.Dropdown( choices=list(vlm_caption.MODEL_REGISTRY.keys()), value=vlm_caption.DEFAULT_MODEL, label="VLM model (ignored in Detector-only mode)", ), ], outputs=[ gr.Image(label="Annotated detections"), gr.Markdown(label="Safety caption"), gr.JSON(label="Raw detections + compliance"), ], title="AiGiS Safety Captioning", description=( "Upload a workplace image. Choose a mode: the full pipeline (YOLO checks PPE " "per person, then a VLM states the issue, the main hazard and the missing " "controls and writes the safety observation), detector only, or VLM only " "(the model judges the image without object detection)." ), ) if __name__ == "__main__": demo.launch()