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