Clement Delteil
commited on
Commit
•
b55565e
1
Parent(s):
fe63512
fix resolution video
Browse files
app.py
CHANGED
@@ -1,8 +1,9 @@
|
|
|
|
|
|
1 |
import streamlit as st
|
2 |
from models.deep_colorization.colorizers import *
|
3 |
import cv2
|
4 |
from PIL import Image
|
5 |
-
import pathlib
|
6 |
import tempfile
|
7 |
import moviepy.editor as mp
|
8 |
import time
|
@@ -35,20 +36,8 @@ def colorize_frame(frame, colorizer) -> np.ndarray:
|
|
35 |
tens_l_orig, tens_l_rs = preprocess_img(frame, HW=(256, 256))
|
36 |
return postprocess_tens(tens_l_orig, colorizer(tens_l_rs).cpu())
|
37 |
|
38 |
-
|
39 |
image = Image.open(r'img/streamlit.png') # Brand logo image (optional)
|
40 |
|
41 |
-
APP_DIR = pathlib.Path(__file__).parent.absolute()
|
42 |
-
|
43 |
-
LOCAL_DIR = APP_DIR / "local_video"
|
44 |
-
LOCAL_DIR.mkdir(exist_ok=True)
|
45 |
-
save_dir = LOCAL_DIR / "output"
|
46 |
-
save_dir.mkdir(exist_ok=True)
|
47 |
-
|
48 |
-
print(APP_DIR)
|
49 |
-
print(LOCAL_DIR)
|
50 |
-
print(save_dir)
|
51 |
-
|
52 |
# Create two columns with different width
|
53 |
col1, col2 = st.columns([0.8, 0.2])
|
54 |
with col1: # To display the header text using css style
|
@@ -72,7 +61,7 @@ uploaded_file = st.file_uploader("", type=['jpg', 'png', 'jpeg', 'mp4'])
|
|
72 |
|
73 |
# Add 'before' and 'after' columns
|
74 |
if uploaded_file is not None:
|
75 |
-
file_extension =
|
76 |
|
77 |
if file_extension in ['jpg', 'png', 'jpeg']:
|
78 |
image = Image.open(uploaded_file)
|
@@ -90,13 +79,13 @@ if uploaded_file is not None:
|
|
90 |
if filter == 'ECCV 16':
|
91 |
colorizer_eccv16 = eccv16(pretrained=True).eval()
|
92 |
img = load_img(uploaded_file)
|
93 |
-
|
94 |
out_img_eccv16 = postprocess_tens(tens_l_orig, colorizer_eccv16(tens_l_rs).cpu())
|
95 |
st.image(out_img_eccv16, width=300)
|
96 |
elif filter == 'SIGGRAPH 17':
|
97 |
colorizer_siggraph17 = siggraph17(pretrained=True).eval()
|
98 |
img = load_img(uploaded_file)
|
99 |
-
|
100 |
out_img_siggraph17 = postprocess_tens(tens_l_orig, colorizer_siggraph17(tens_l_rs).cpu())
|
101 |
st.image(out_img_siggraph17, width=300)
|
102 |
else:
|
@@ -156,13 +145,10 @@ if uploaded_file is not None:
|
|
156 |
progress_bar.empty()
|
157 |
|
158 |
with st.spinner("Merging frames to video..."):
|
159 |
-
print("finished")
|
160 |
frame_size = output_frames[0].shape[:2]
|
161 |
-
print(frame_size)
|
162 |
output_filename = "output.mp4"
|
163 |
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # Codec for MP4 video
|
164 |
-
|
165 |
-
out = cv2.VideoWriter(output_filename, fourcc, fps, (3840, 2160))
|
166 |
|
167 |
# Display the colorized video using st.video
|
168 |
for frame in output_frames:
|
|
|
1 |
+
import os.path
|
2 |
+
|
3 |
import streamlit as st
|
4 |
from models.deep_colorization.colorizers import *
|
5 |
import cv2
|
6 |
from PIL import Image
|
|
|
7 |
import tempfile
|
8 |
import moviepy.editor as mp
|
9 |
import time
|
|
|
36 |
tens_l_orig, tens_l_rs = preprocess_img(frame, HW=(256, 256))
|
37 |
return postprocess_tens(tens_l_orig, colorizer(tens_l_rs).cpu())
|
38 |
|
|
|
39 |
image = Image.open(r'img/streamlit.png') # Brand logo image (optional)
|
40 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
41 |
# Create two columns with different width
|
42 |
col1, col2 = st.columns([0.8, 0.2])
|
43 |
with col1: # To display the header text using css style
|
|
|
61 |
|
62 |
# Add 'before' and 'after' columns
|
63 |
if uploaded_file is not None:
|
64 |
+
file_extension = os.path.splitext(uploaded_file.name)[1].lower()
|
65 |
|
66 |
if file_extension in ['jpg', 'png', 'jpeg']:
|
67 |
image = Image.open(uploaded_file)
|
|
|
79 |
if filter == 'ECCV 16':
|
80 |
colorizer_eccv16 = eccv16(pretrained=True).eval()
|
81 |
img = load_img(uploaded_file)
|
82 |
+
tens_l_orig, tens_l_rs = preprocess_img(img, HW=(256, 256))
|
83 |
out_img_eccv16 = postprocess_tens(tens_l_orig, colorizer_eccv16(tens_l_rs).cpu())
|
84 |
st.image(out_img_eccv16, width=300)
|
85 |
elif filter == 'SIGGRAPH 17':
|
86 |
colorizer_siggraph17 = siggraph17(pretrained=True).eval()
|
87 |
img = load_img(uploaded_file)
|
88 |
+
tens_l_orig, tens_l_rs = preprocess_img(img, HW=(256, 256))
|
89 |
out_img_siggraph17 = postprocess_tens(tens_l_orig, colorizer_siggraph17(tens_l_rs).cpu())
|
90 |
st.image(out_img_siggraph17, width=300)
|
91 |
else:
|
|
|
145 |
progress_bar.empty()
|
146 |
|
147 |
with st.spinner("Merging frames to video..."):
|
|
|
148 |
frame_size = output_frames[0].shape[:2]
|
|
|
149 |
output_filename = "output.mp4"
|
150 |
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # Codec for MP4 video
|
151 |
+
out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_size[1], frame_size[0]))
|
|
|
152 |
|
153 |
# Display the colorized video using st.video
|
154 |
for frame in output_frames:
|