Update app.py

app.py

@@ -1,39 +1,41 @@
 import torch
-import torch.nn as nn
 import gradio as gr
 import cv2
 import numpy as np
 
-# Define your model architecture
-class YourModelArchitecture(nn.Module):
+# Chemin vers le modèle
+MODEL_PATH = 'ColorizeVideo_gen.pth'
+
+# Définir l'architecture de votre modèle ici
+class YourModelArchitecture(torch.nn.Module):
     def __init__(self):
         super(YourModelArchitecture, self).__init__()
-        # Define the layers of your model here
-        # Example: self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
-    
-    def forward(self, x):
-        # Define the forward pass logic
-        return x
+        # Définissez les couches de votre modèle
 
-# Path to the model weights
-MODEL_PATH = 'ColorizeVideo_gen.pth'
+    def forward(self, x):
+        # Définissez la logique de passage avant de votre modèle
+        return x  # Modifiez ceci selon votre modèle
 
-# Load the model function
+# Charger le modèle
 def load_model(model_path):
-    model = YourModelArchitecture()  # Initialize the model architecture
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))  # Load the model weights
-    model.eval()  # Set the model to evaluation mode
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))  # Charger le checkpoint
+    model = YourModelArchitecture()  # Initialiser l'architecture du modèle
+    
+    # Charger uniquement les poids du modèle à partir du checkpoint
+    model.load_state_dict(checkpoint['model'])  
+    
+    model.eval()  # Met le modèle en mode évaluation
     return model
 
-# Preprocess the frame before passing it to the model
+# Prétraitement de l'image
 def preprocess_frame(frame):
-    # Resize and normalize the image
-    frame = cv2.resize(frame, (224, 224))  # Resize to model input size
-    frame = frame / 255.0  # Normalize the pixel values to [0, 1]
-    input_tensor = torch.from_numpy(frame.astype(np.float32)).permute(2, 0, 1)  # Convert to tensor and change the dimension order
-    return input_tensor.unsqueeze(0)  # Add batch dimension
+    # 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
 
-# Process the video, frame by frame
+# Traitement de la vidéo
 def process_video(model, video_path):
     cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
@@ -45,36 +47,36 @@ def process_video(model, video_path):
         if not ret:
             break
         
-        # Preprocess the frame
+        # Prétraiter le cadre
         input_tensor = preprocess_frame(frame)
 
-        # Make predictions with the model
+        # Faire des prédictions
         with torch.no_grad():
             predictions = model(input_tensor)
         
-        # Convert the predictions back to an image format
+        # Traiter les prédictions et convertir en image
         output_frame = (predictions.squeeze().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
         
-        # Write the processed frame to the output video
+        # Écrire le cadre traité dans la sortie
         out.write(output_frame)
 
     cap.release()
     out.release()
     return output_path
 
-# Gradio interface function
+# Interface Gradio
 def colorize_video(video):
     model = load_model(MODEL_PATH)
-    output_video_path = process_video(model, video.name)  # Use the video file name to read the video
+    output_video_path = process_video(model, video.name)  # Utiliser le nom pour lire la vidéo
     return output_video_path
 
-# Configure the Gradio interface
+# Configuration de l'interface Gradio
 iface = gr.Interface(
     fn=colorize_video,
-    inputs=gr.Video(label="Upload a black and white video"),
-    outputs=gr.Video(label="Colorized Video"),
-    title="Video Colorization",
-    description="Upload a black and white video to colorize it using the model."
+    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__':