Spaces:
Paused
Paused
| """ | |
| AiGiS Safety Captioning — Hugging Face Space (Gradio UI). | |
| Same pipeline as the FastAPI service in ../src/api.py, wrapped in a simple | |
| Gradio interface for HF Spaces: | |
| image in -> YOLO (detector.py) -> structured detections + compliance | |
| -> Claude (vlm_caption.py) -> safety observation | |
| -> annotated image + caption shown in the UI | |
| The YOLO weights are pulled once at startup from the Hugging Face model repo | |
| below, so no local .pt file is needed in this Space. | |
| Every Gradio app also exposes an automatic API (see the "Use via API" link at | |
| the bottom of the Space page), so this doubles as a callable endpoint too. | |
| Requires the ANTHROPIC_API_KEY secret to be set in the Space settings. | |
| """ | |
| 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", "omarkashif/hockey-yolo-models") | |
| HF_FILENAME = os.environ.get("AIGIS_HF_FILENAME", "aigis/yolov8s.pt") | |
| DATA_YAML = os.environ.get("AIGIS_DATA_YAML", "sh17.yaml") | |
| CONF = float(os.environ.get("AIGIS_CONF", "0.25")) | |
| # Load the model on CPU at startup. Actual GPU attachment (ZeroGPU) only | |
| # happens inside the @spaces.GPU-decorated function below, at inference time. | |
| weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME) | |
| detector = SafetyDetector(weights_path, DATA_YAML, conf=CONF) | |
| def run(image: np.ndarray, skip_vlm: bool): | |
| if image is None: | |
| return None, "Upload an image first.", {} | |
| # Gradio gives RGB; Ultralytics/OpenCV (and our detector) work in BGR. | |
| 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 skip_vlm: | |
| return annotated_rgb, "_(VLM captioning skipped)_", det_dict | |
| if "ANTHROPIC_API_KEY" not in os.environ: | |
| return annotated_rgb, "**Error:** ANTHROPIC_API_KEY is not set for this Space.", det_dict | |
| ok, buf = cv2.imencode(".jpg", frame_bgr) | |
| vlm_result = vlm_caption.caption_bytes(buf.tobytes(), "image/jpeg", det_dict) | |
| caption_md = ( | |
| f"**Caption:** {vlm_result.get('caption', '(parse error)')}\n\n" | |
| f"**Safety observation:** {vlm_result.get('safety_observation', '')}\n\n" | |
| f"**Severity:** {vlm_result.get('severity', 'n/a')}" | |
| ) | |
| return annotated_rgb, caption_md, det_dict | |
| demo = gr.Interface( | |
| fn=run, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload image"), | |
| gr.Checkbox(label="Skip VLM caption (detector only, no API cost)"), | |
| ], | |
| 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. YOLO detects PPE compliance deterministically; " | |
| "Claude turns those findings into a grounded safety observation." | |
| ), | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |