Update app.py

app.py

@@ -1,71 +1,80 @@
 import torch
+import torch.nn as nn
 import gradio as gr
 import cv2
 import numpy as np
 
-# Chemin vers le modèle
+# Define your model architecture
+class YourModelArchitecture(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
+
+# Path to the model weights
 MODEL_PATH = 'ColorizeVideo_gen.pth'
 
-# Charger le modèle
+# Load the model function
 def load_model(model_path):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = torch.load(model_path, map_location=device)
-    model.eval()
+    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
     return model
 
-# Prétraitement de l'image
+# Preprocess the frame before passing it to the model
 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)
-    return input_tensor.unsqueeze(0)
+    # 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
 
-# Traitement de la vidéo
+# Process the video, frame by frame
 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, (width, height))
+    out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         
-        # Prétraiter le cadre
+        # Preprocess the frame
         input_tensor = preprocess_frame(frame)
 
-        # Faire des prédictions
+        # Make predictions with the model
         with torch.no_grad():
             predictions = model(input_tensor)
         
-        # Convertir en image
+        # Convert the predictions back to an image format
         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
+        # Write the processed frame to the output video
         out.write(output_frame)
 
     cap.release()
     out.release()
     return output_path
 
-# Interface Gradio
+# Gradio interface function
 def colorize_video(video):
     model = load_model(MODEL_PATH)
-    output_video_path = process_video(model, video.name)
+    output_video_path = process_video(model, video.name)  # Use the video file name to read the video
     return output_video_path
 
-# Configuration de l'interface Gradio
+# Configure the Gradio interface
 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."
+    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."
 )
 
 if __name__ == '__main__':