Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| 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() |