Create Integ.py

Integ.py

@@ -0,0 +1,63 @@
+import os
+import torch
+import cv2
+from fastapi import FastAPI, File, UploadFile
+from ultralytics import YOLO
+import io
+from pydantic import BaseModel
+import json
+
+# Initialize the Hugging Face app
+app = FastAPI()
+
+# Load YOLO model
+model = YOLO('yolov8s.pt')  # Replace with your custom model if needed
+
+class Violation(BaseModel):
+    frame: int
+    violation: str
+    confidence: float
+    timestamp: str
+
+# Inference endpoint to process video
+@app.post("/process_video/")
+async def process_video(file: UploadFile = File(...)):
+    # Read video file from user upload
+    video_data = await file.read()
+    video_path = f"temp_{file.filename}"
+    with open(video_path, "wb") as f:
+        f.write(video_data)
+
+    # Process the video using YOLO model
+    results = process_video_with_yolo(video_path)
+
+    # Clean up
+    os.remove(video_path)
+
+    return {"violations": results}
+
+# Function to process video with YOLO model
+def process_video_with_yolo(video_path):
+    video = cv2.VideoCapture(video_path)
+    violations = []
+    frame_count = 0
+
+    while True:
+        ret, frame = video.read()
+        if not ret:
+            break
+        # Run inference on the frame
+        results = model(frame)
+        for result in results:
+            for box in result.boxes:
+                violation = {
+                    "frame": frame_count,
+                    "violation": "no_helmet",  # Example violation (this can be dynamic based on the class)
+                    "confidence": float(box.conf),
+                    "timestamp": str(frame_count / 30)  # Timestamp (frame rate assumption)
+                }
+                violations.append(violation)
+        frame_count += 1
+
+    video.release()
+    return violations