Spaces:
Sleeping
Sleeping
File size: 1,727 Bytes
b69e47e |
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 |
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
|