import torch
import gradio as gr
import cv2
import numpy as np

# Chemin vers le modèle
MODEL_PATH = 'ColorizeVideo_gen.pth'

# Charger le modèle
def load_model(model_path):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load(model_path, map_location=device)
    model.eval()
    return model

# Prétraitement de l'image
def preprocess_frame(frame):
    # Redimensionner et normaliser
    frame = cv2.resize(frame, (224, 224))  # Ajustez la taille si nécessaire
    frame = frame / 255.0  # Normaliser
    input_tensor = torch.from_numpy(frame.astype(np.float32)).permute(2, 0, 1)
    return input_tensor.unsqueeze(0)

# Traitement de la vidéo
def process_video(model, video_path):
    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    output_path = "output_video.mp4"
    out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Prétraiter le cadre
        input_tensor = preprocess_frame(frame)

        # Faire des prédictions
        with torch.no_grad():
            predictions = model(input_tensor)
        
        # Convertir en image
        output_frame = (predictions.squeeze().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        output_frame = cv2.resize(output_frame, (frame.shape[1], frame.shape[0]))  # Rétablir la taille originale
        
        # Écrire le cadre traité dans la sortie
        out.write(output_frame)

    cap.release()
    out.release()
    return output_path

# Interface Gradio
def colorize_video(video):
    model = load_model(MODEL_PATH)
    output_video_path = process_video(model, video.name)
    return output_video_path

# Configuration de l'interface Gradio
iface = gr.Interface(
    fn=colorize_video,
    inputs=gr.Video(label="Téléchargez une vidéo"),
    outputs=gr.Video(label="Vidéo colorisée"),
    title="Colorisation de Vidéos",
    description="Chargez une vidéo en noir et blanc et utilisez le modèle de colorisation pour obtenir une vidéo colorisée."
)

if __name__ == '__main__':
    iface.launch()