import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
# Charger le modèle YOLOv8 pré-entraîné
model = YOLO("yolov8n.pt")

# Fonction pour la détection sur image
def detect_objects_image(img):
    results = model(img)  # Détection
    annotated_frame = results[0].plot()  # Annoter les résultats
    return annotated_frame


def detect_objects_video(video_path):
    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out_path = "annotated_video.mp4"
    out = cv2.VideoWriter(out_path, fourcc, 20.0, (int(cap.get(3)), int(cap.get(4))))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        result = model(frame)
        annotated = result[0].plot()
        out.write(annotated)

    cap.release()
    out.release()
    return out_path

demo = gr.Blocks(theme='JohnSmith9982/small_and_pretty')

#Interface Gradio

image_input = gr.Image(type="numpy", label="Image à analyser")
video_input = gr.Image(label="Video à analyser")

image_output = gr.Image(type="numpy", label="Image annotée")
video_output = gr.Image(label="Video annotée")

interface1 = gr.Interface(fn=detect_objects_image, inputs=image_input, outputs=image_output, title="Détection sur Image")
interface2 = gr.Interface(fn=detect_objects_video, inputs=video_input, outputs=video_output, title="Détection sur Video")

with demo: 
  gr.TabbedInterface(
      [interface1, interface2],
      ['Image', 'Video']
  )
demo.launch()