File size: 3,999 Bytes
87e2c95 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """
Gradio App for PPE Compliance Detection
Run: gradio gradio_app.py
"""
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
from PIL import Image
import cv2
import numpy as np
MODEL_ID = "baskarmother/yolov8-ppe-construction"
CLASS_NAMES = [
'barricade', 'dumpster', 'excavators', 'gloves', 'hardhat', 'mask',
'no-hardhat', 'no-mask', 'no-safety vest', 'person', 'safety net',
'safety shoes', 'safety vest', 'dump truck', 'mini-van', 'truck', 'wheel loader'
]
def get_color(label):
if label in ["hardhat", "mask", "safety vest", "gloves", "safety shoes"]:
return (0, 255, 0) # Green = compliant PPE
elif label in ["no-hardhat", "no-mask", "no-safety vest"]:
return (0, 0, 255) # Red = violation
else:
return (255, 0, 0) # Blue = other
# Load model once
print("Loading model from Hub...")
weights_path = hf_hub_download(MODEL_ID, "best.pt")
model = YOLO(weights_path)
print("Model loaded!")
def detect_ppe(image, conf_threshold=0.25):
results = model(image, conf=conf_threshold)
result = results[0]
# Convert PIL to numpy array
img = np.array(image)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
detections = []
compliance = {"hardhat": "NOT DETECTED", "mask": "NOT DETECTED", "safety_vest": "NOT DETECTED"}
detected_classes = set()
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
cls_id = int(box.cls[0])
conf = float(box.conf[0])
label = CLASS_NAMES[cls_id]
detected_classes.add(label)
color = get_color(label)
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
text = f"{label} {conf:.2f}"
cv2.putText(img, text, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
detections.append(f"{label}: {conf:.2f}")
# Compliance check
if "hardhat" in detected_classes:
compliance["hardhat"] = "✅ COMPLIANT"
if "no-hardhat" in detected_classes:
compliance["hardhat"] = "❌ VIOLATION"
if "mask" in detected_classes:
compliance["mask"] = "✅ COMPLIANT"
if "no-mask" in detected_classes:
compliance["mask"] = "❌ VIOLATION"
if "safety vest" in detected_classes:
compliance["safety_vest"] = "✅ COMPLIANT"
if "no-safety vest" in detected_classes:
compliance["safety_vest"] = "❌ VIOLATION"
# Convert back to RGB
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
output_image = Image.fromarray(img_rgb)
report = f"""### PPE Compliance Report
- **Hard Hat**: {compliance['hardhat']}
- **Mask**: {compliance['mask']}
- **Safety Vest**: {compliance['safety_vest']}
### Detections ({len(detections)} objects)
{chr(10).join(['- ' + d for d in detections]) if detections else 'No objects detected'}
"""
return output_image, report
with gr.Blocks(title="PPE Compliance Detection") as demo:
gr.Markdown("# 🦺 PPE Compliance Detection for Construction Sites")
gr.Markdown("Upload a construction site image to detect Personal Protective Equipment (PPE) compliance.")
with gr.Row():
with gr.Column():
input_image = gr.Image(type="pil", label="Upload Image")
conf_slider = gr.Slider(0.1, 0.9, value=0.25, step=0.05, label="Confidence Threshold")
detect_btn = gr.Button("Detect PPE", variant="primary")
with gr.Column():
output_image = gr.Image(label="Detected Objects")
output_text = gr.Markdown(label="Compliance Report")
detect_btn.click(fn=detect_ppe, inputs=[input_image, conf_slider], outputs=[output_image, output_text])
gr.Markdown("""
## Classes Detected
- ✅ **Compliant PPE**: hardhat, mask, safety vest, gloves, safety shoes
- ❌ **Violations**: no-hardhat, no-mask, no-safety vest
- 🔵 **Other**: person, equipment, vehicles, barricades
""")
if __name__ == "__main__":
demo.launch()
|