fix resolution video

app.py

@@ -1,8 +1,9 @@
+import os.path
+
 import streamlit as st
 from models.deep_colorization.colorizers import *
 import cv2
 from PIL import Image
-import pathlib
 import tempfile
 import moviepy.editor as mp
 import time
@@ -35,20 +36,8 @@ def colorize_frame(frame, colorizer) -> np.ndarray:
     tens_l_orig, tens_l_rs = preprocess_img(frame, HW=(256, 256))
     return postprocess_tens(tens_l_orig, colorizer(tens_l_rs).cpu())
 
-
 image = Image.open(r'img/streamlit.png')  # Brand logo image (optional)
 
-APP_DIR = pathlib.Path(__file__).parent.absolute()
-
-LOCAL_DIR = APP_DIR / "local_video"
-LOCAL_DIR.mkdir(exist_ok=True)
-save_dir = LOCAL_DIR / "output"
-save_dir.mkdir(exist_ok=True)
-
-print(APP_DIR)
-print(LOCAL_DIR)
-print(save_dir)
-
 # Create two columns with different width
 col1, col2 = st.columns([0.8, 0.2])
 with col1:  # To display the header text using css style
@@ -72,7 +61,7 @@ uploaded_file = st.file_uploader("", type=['jpg', 'png', 'jpeg', 'mp4'])
 
 # Add 'before' and 'after' columns
 if uploaded_file is not None:
-    file_extension = uploaded_file.name.split('.')[1].lower()
+    file_extension = os.path.splitext(uploaded_file.name)[1].lower()
 
     if file_extension in ['jpg', 'png', 'jpeg']:
         image = Image.open(uploaded_file)
@@ -90,13 +79,13 @@ if uploaded_file is not None:
             if filter == 'ECCV 16':
                 colorizer_eccv16 = eccv16(pretrained=True).eval()
                 img = load_img(uploaded_file)
-                (tens_l_orig, tens_l_rs) = preprocess_img(img, HW=(256, 256))
+                tens_l_orig, tens_l_rs = preprocess_img(img, HW=(256, 256))
                 out_img_eccv16 = postprocess_tens(tens_l_orig, colorizer_eccv16(tens_l_rs).cpu())
                 st.image(out_img_eccv16, width=300)
             elif filter == 'SIGGRAPH 17':
                 colorizer_siggraph17 = siggraph17(pretrained=True).eval()
                 img = load_img(uploaded_file)
-                (tens_l_orig, tens_l_rs) = preprocess_img(img, HW=(256, 256))
+                tens_l_orig, tens_l_rs = preprocess_img(img, HW=(256, 256))
                 out_img_siggraph17 = postprocess_tens(tens_l_orig, colorizer_siggraph17(tens_l_rs).cpu())
                 st.image(out_img_siggraph17, width=300)
             else:
@@ -156,13 +145,10 @@ if uploaded_file is not None:
                             progress_bar.empty()
 
                 with st.spinner("Merging frames to video..."):
-                    print("finished")
                     frame_size = output_frames[0].shape[:2]
-                    print(frame_size)
                     output_filename = "output.mp4"
                     fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for MP4 video
-                    print(fps)
-                    out = cv2.VideoWriter(output_filename, fourcc, fps, (3840, 2160))
+                    out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_size[1], frame_size[0]))
 
                     # Display the colorized video using st.video
                     for frame in output_frames: