app.py

import torch
import cv2
import numpy as np

# Chemins vers le modèle et la vidéo à traiter
MODEL_PATH = 'ColorizeVideo_gen.pth'
VIDEO_PATH = '119195-716970703_small.mp4'  # Nom de la vidéo exemple
OUTPUT_VIDEO_PATH = 'output_video.mp4'

# Charger le modèle
def load_model(model_path):
    model = torch.load(model_path, map_location=torch.device('cpu'))  # Charger sur le CPU
    model.eval()  # Met le modèle en mode évaluation
    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)  # Convertir en format Tensor
    return input_tensor.unsqueeze(0)  # Ajouter une dimension de lot

# Traitement de la vidéo
def process_video(model, video_path, output_path):
    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))

    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)
        
        # Traiter les prédictions et convertir en image
        output_frame = (predictions.squeeze().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        
        # Écrire le cadre traité dans la sortie
        out.write(output_frame)

    cap.release()
    out.release()

# Fonction principale
if __name__ == '__main__':
    model = load_model(MODEL_PATH)
    process_video(model, VIDEO_PATH, OUTPUT_VIDEO_PATH)
    print(f"Traitement de la vidéo terminé. Résultats enregistrés dans {OUTPUT_VIDEO_PATH}.")