solfedge's picture
Upload 14 files
6340002 verified
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)