Update app.py

app.py

@@ -1,11 +1,10 @@
 import torch
+import gradio as gr
 import cv2
 import numpy as np
 
-# Chemins vers le modèle et la vidéo à traiter
+# Chemin vers le modèle
 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):
@@ -22,9 +21,10 @@ def preprocess_frame(frame):
     return input_tensor.unsqueeze(0)  # Ajouter une dimension de lot
 
 # Traitement de la vidéo
-def process_video(model, video_path, output_path):
+def process_video(model, video_path):
     cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    output_path = "output_video.mp4"
     out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))
 
     while cap.isOpened():
@@ -47,9 +47,22 @@ def process_video(model, video_path, output_path):
 
     cap.release()
     out.release()
+    return output_path
 
-# Fonction principale
-if __name__ == '__main__':
+# Interface Gradio
+def colorize_video(video):
     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}.")
+    output_video_path = process_video(model, video.name)  # Utiliser le nom pour lire la vidéo
+    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()