Update app.py

app.py

@@ -1,20 +1,55 @@
-from deoldify.visualize import get_video_colorizer
 import torch
+import cv2
+import numpy as np
 
-# Load the video colorization model
-def load_video_model():
-    model_path = 'ColorizeVideo_gen.pth'  # Path to the model in the root directory
-    video_colorizer = get_video_colorizer()
+# 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'
 
-    # Load the model's state from the .pth file
-    state = torch.load(model_path, map_location=torch.device('cpu'))  # Adjust if using GPU
-    video_colorizer.learn.model.load_state_dict(state)
+# 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
 
-    return video_colorizer
+# 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
 
-# Example usage of the video colorizer
-video_colorizer = load_video_model()
+# 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))))
 
-# You can now use the colorizer to colorize a video
-video_path = 'your_video.mp4'
-colorized_video = video_colorizer.colorize_from_file_name(video_path, render_factor=35)
+    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}.")