Update app.py

app.py

@@ -1,54 +1,18 @@
-import os
-import sys
-import time
-import gradio as gr
 import torch
-import cv2
-
-# Cloner le repo si non présent
-if not os.path.exists('TCVC-Temporally-Consistent-Video-Colorization'):
-    os.system('git clone https://github.com/lyh-18/TCVC-Temporally-Consistent-Video-Colorization.git')
-
-# Vérifier si le dépôt est bien cloné
-while not os.path.exists('TCVC-Temporally-Consistent-Video-Colorization'):
-    time.sleep(1)  # Attendre que le clonage soit terminé
-
-# Ajouter le dépôt cloné au chemin Python
-sys.path.append('TCVC-Temporally-Consistent-Video-Colorization')
+from deoldify.visualize import get_image_colorizer, get_video_colorizer
+import gradio as gr
 
-# Importer le modèle une fois le dépôt cloné
-from TCVC_Temporally_Consistent_Video_Colorization.models.TCVC import TCVC_Model
+# Clone le repo et télécharge les dépendances
+if not os.path.exists('DeOldify'):
+    os.system('git clone https://github.com/jantic/DeOldify.git')
+    os.system('pip install -r DeOldify/requirements-colab.txt')
 
-# Charger le modèle pré-entraîné
-def load_model():
-    model = TCVC_Model()
-    model.load_state_dict(torch.load("path_to_pretrained_model.pth"))  # Charge le modèle pré-entraîné
-    model.eval()  # Mode évaluation
-    return model
+# Charger le modèle DeOldify pré-entraîné
+video_colorizer = get_video_colorizer()
 
-# Fonction de colorisation de la vidéo
+# Fonction pour coloriser une vidéo
 def colorize_video(video_path):
-    model = load_model()
-
-    # Lire la vidéo
-    cap = cv2.VideoCapture(video_path)
-    frame_width = int(cap.get(3))
-    frame_height = int(cap.get(4))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-
-    output_path = "colorized_output.mp4"
-    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
-
-    while(cap.isOpened()):
-        ret, frame = cap.read()
-        if not ret:
-            break
-        # Convertir l'image en noir et blanc en couleur
-        color_frame = model.colorize_frame(frame)
-        out.write(color_frame)
-
-    cap.release()
-    out.release()
+    output_path = video_colorizer.colorize_from_file_name(video_path)
     return output_path
 
 # Interface Gradio
@@ -56,12 +20,12 @@ def process_video(video):
     colorized_video = colorize_video(video.name)
     return colorized_video
 
-# Créer l'interface Gradio
+# Interface Gradio pour uploader la vidéo
 interface = gr.Interface(fn=process_video, 
-                         inputs=gr.Video(label="Télécharger une vidéo en noir et blanc"),
+                         inputs=gr.Video(label="Télécharger une vidéo en noir et blanc"), 
                          outputs=gr.Video(label="Vidéo colorisée"),
-                         title="Colorisation vidéo",
-                         description="Téléchargez une vidéo en noir et blanc et obtenez une version colorisée.")
+                         title="Colorisation de Vidéo avec DeOldify",
+                         description="Téléchargez une vidéo en noir et blanc pour la coloriser.")
 
 if __name__ == "__main__":
     interface.launch()