Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from models.TCVC import TCVC_Model # Assurez-vous que le chemin est correct après avoir cloné le repo | |
| import cv2 | |
| import os | |
| # 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 | |
| # Fonction de colorisation de la 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() | |
| return output_path | |
| # Interface Gradio | |
| def process_video(video): | |
| colorized_video = colorize_video(video.name) | |
| return colorized_video | |
| # Créer l'interface Gradio | |
| interface = gr.Interface(fn=process_video, | |
| 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.") | |
| # Lancer l'interface | |
| if __name__ == "__main__": | |
| interface.launch() | |