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

import tempfile
# Fonction pour la détection sur vidéo
def detect_objects_video(video):
    # Si l'entrée est une chaîne, utiliser telle quelle. Sinon, utiliser .name (cas Gradio)
    video_path = video.name if hasattr(video, 'name') else video

    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    cap = cv2.VideoCapture(video_path)

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height))

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        results = model(frame)
        annotated_frame = results[0].plot()
        out.write(annotated_frame)

    cap.release()
    out.release()

    return temp_output.name


demo = gr.Blocks(theme='shivi/calm_seafoam')

#Interface Gradio
image_input = gr.Image(type="numpy", label="Image à analyser")
video_input = gr.Video(label="Vidéo à analyser")
image_output = gr.Image(type="numpy", label="Image annotée")
video_output = gr.Video(label="Vidéo annotée")

interface1 = gr.Interface(fn=detect_objects_image, 
                          inputs=image_input, 
                          outputs=image_output, 
                          title="Détection sur Image", 
                          description="""
    Cette interface permet de détecter automatiquement les objets présents sur une image.

    Le modèle YOLOv8 est utilisé pour effectuer une détection rapide et précise.

    Téléversez une image, et les objets détectés seront annotés visuellement.
    """)
interface2 = gr.Interface(fn=detect_objects_video, 
                          inputs=video_input, 
                          outputs=video_output, 
                          title="Détection sur Vidéo", 
                          description="""
    Cette interface permet d’analyser une vidéo et d’y détecter les objets image par image.

    Le modèle YOLOv8 est appliqué à chaque image de la vidéo pour générer une version annotée.

    Téléversez une vidéo, puis récupérez la vidéo traitée avec les objets détectés en surbrillance.
    """
)
with demo:
  gr.TabbedInterface(
      [interface1 ,interface2],
      ['Image', 'Video']
  )

demo.launch()