Update app.py

app.py

@@ -8,8 +8,9 @@ MODEL_PATH = 'ColorizeVideo_gen.pth'
 
 # 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
+    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
@@ -17,15 +18,17 @@ 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
+    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, (int(cap.get(3)), int(cap.get(4))))
+    out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))
 
     while cap.isOpened():
         ret, frame = cap.read()
@@ -39,8 +42,9 @@ def process_video(model, video_path):
         with torch.no_grad():
             predictions = model(input_tensor)
         
-        # Traiter les prédictions et convertir en image
+        # 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)
@@ -52,7 +56,7 @@ def process_video(model, video_path):
 # Interface Gradio
 def colorize_video(video):
     model = load_model(MODEL_PATH)
-    output_video_path = process_video(model, video.name)  # Utiliser le nom pour lire la vidéo
+    output_video_path = process_video(model, video.name)
     return output_video_path
 
 # Configuration de l'interface Gradio