Add inference script for PPE detection

inference.py

@@ -0,0 +1,115 @@
+"""
+PPE Compliance Detection - Inference Script
+Usage with trained YOLOv8 model
+"""
+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"
+
+# PPE compliance pairs: (required_ppe, violation_class)
+PPE_PAIRS = {
+    "hardhat": ("hardhat", "no-hardhat"),
+    "mask": ("mask", "no-mask"),
+    "safety_vest": ("safety vest", "no-safety vest"),
+}
+
+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 load_model():
+    """Load model from HuggingFace Hub."""
+    weights_path = hf_hub_download(MODEL_ID, "best.pt")
+    model = YOLO(weights_path)
+    return model
+
+def detect(image_path, conf_threshold=0.25):
+    """Run detection on an image."""
+    model = load_model()
+    results = model(image_path, conf=conf_threshold)
+    return results[0]
+
+def check_compliance(result):
+    """Check PPE compliance from detection results."""
+    boxes = result.boxes
+    detected_classes = set()
+    
+    for cls in boxes.cls:
+        detected_classes.add(CLASS_NAMES[int(cls)])
+    
+    compliance_report = {}
+    for ppe_name, (required, violation) in PPE_PAIRS.items():
+        has_required = required in detected_classes
+        has_violation = violation in detected_classes
+        
+        if has_violation:
+            status = "VIOLATION"
+        elif has_required:
+            status = "COMPLIANT"
+        else:
+            status = "NOT DETECTED"
+        
+        compliance_report[ppe_name] = {
+            "status": status,
+            "required_detected": has_required,
+            "violation_detected": has_violation
+        }
+    
+    return compliance_report
+
+def draw_detections(image_path, result, save_path="output.jpg"):
+    """Draw bounding boxes and labels on image."""
+    img = cv2.imread(image_path)
+    boxes = result.boxes
+    
+    for box in 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]
+        
+        # Color: green for compliant PPE, red for violations, blue for others
+        if label in ["hardhat", "mask", "safety vest", "gloves", "safety shoes"]:
+            color = (0, 255, 0)
+        elif label in ["no-hardhat", "no-mask", "no-safety vest"]:
+            color = (0, 0, 255)
+        else:
+            color = (255, 0, 0)
+        
+        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+        text = f"{label} {conf:.2f}"
+        cv2.putText(img, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    
+    cv2.imwrite(save_path, img)
+    return save_path
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", required=True, help="Path to input image")
+    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
+    parser.add_argument("--output", default="output.jpg", help="Output image path")
+    args = parser.parse_args()
+    
+    print("Loading model...")
+    result = detect(args.image, args.conf)
+    
+    print(f"Detected {len(result.boxes)} objects")
+    for box in result.boxes:
+        cls_id = int(box.cls[0])
+        conf = float(box.conf[0])
+        print(f"  - {CLASS_NAMES[cls_id]}: {conf:.3f}")
+    
+    print("\n--- PPE Compliance Report ---")
+    report = check_compliance(result)
+    for ppe, info in report.items():
+        print(f"  {ppe}: {info['status']}")
+    
+    draw_detections(args.image, result, args.output)
+    print(f"\nAnnotated image saved to {args.output}")