Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from PIL import Image | |
| import os | |
| # Import modules | |
| from detect_anomaly import detect_hazard | |
| from generate_hazards import generate_unsafe_image | |
| from explain_alert import explain_hazard | |
| from detect_helmet import detect_helmet_in_image | |
| last_result = {"anomalous": False, "error": 0.0} | |
| def analyze_image(img): | |
| global last_result | |
| if img is None: | |
| return "Please upload an image", None, None | |
| pil_img = Image.fromarray(img).convert("RGB") | |
| is_anomalous, error, recon = detect_hazard(pil_img) | |
| last_result = {"anomalous": is_anomalous, "error": error} | |
| alert = "HAZARD DETECTED: No helmet or danger zone!" if is_anomalous else "No hazard detected." | |
| explanation = explain_hazard( | |
| helmet="No" if is_anomalous else "Yes", | |
| zone="Danger" if is_anomalous else "Safe", | |
| lighting="Normal" | |
| ) | |
| return alert, recon, explanation | |
| def create_synthetic(): | |
| img = generate_unsafe_image() | |
| return img | |
| #Gradio Interface | |
| with gr.Blocks(title="AI Safety Inspector") as demo: | |
| gr.Markdown("# AI Safety Inspector\nDetects missing helmets, danger zones using **Unsupervised + Gen AI**") | |
| with gr.Tabs(): | |
| with gr.Tab("Analyze Image"): | |
| with gr.Row(): | |
| input_img = gr.Image(label="Upload Site Photo") | |
| output_recon = gr.Image(label="Reconstructed (Autoencoder)") | |
| output_alert = gr.Label(label="Status") | |
| output_explain = gr.Textbox(label="AI Safety Officer Says") | |
| btn = gr.Button("Analyze") | |
| btn.click(analyze_image, inputs=input_img, outputs=[output_alert, output_recon, output_explain]) | |
| with gr.Tab("Helmet Detection (YOLOv8)"): | |
| gr.Markdown("Detects workers and checks if they are wearing helmets using YOLOv8.") | |
| with gr.Row(): | |
| yolo_input = gr.Image(label="Upload Image") | |
| yolo_output_img = gr.Image(label="Detected Helmets") | |
| yolo_output_count = gr.Number(label="Workers Without Helmet") | |
| yolo_output_labels = gr.Textbox(label="Detections") | |
| yolo_btn = gr.Button("Run Helmet Detection") | |
| yolo_btn.click( | |
| detect_helmet_in_image, | |
| inputs=yolo_input, | |
| outputs=[yolo_output_img, yolo_output_count, yolo_output_labels] | |
| ) | |
| with gr.Tab("Generate Synthetic Hazard (Gen AI)"): | |
| gen_output = gr.Image(label="Generated Unsafe Scenario") | |
| gen_btn = gr.Button("Generate No-Helmet Danger Scene") | |
| gen_btn.click(create_synthetic, outputs=gen_output) | |
| with gr.Tab("ℹ About"): | |
| gr.Markdown(""" | |
| ### How It Works | |
| - Uses **autoencoder** to detect anomalies (no labels needed!) | |
| - **YOLOv8** detects helmets with bounding boxes | |
| - **Stable Diffusion** generates synthetic unsafe images | |
| - **TinyLlama** explains alerts in natural language | |
| - Runs on **Google Colab** | |
| """) | |
| #Launch | |
| demo.launch(share=True) | |