Delete app.py
Browse files
app.py
DELETED
@@ -1,247 +0,0 @@
|
|
1 |
-
import cv2
|
2 |
-
import numpy as np
|
3 |
-
import os
|
4 |
-
import tempfile
|
5 |
-
from tqdm import tqdm
|
6 |
-
import gradio as gr
|
7 |
-
import ffmpeg
|
8 |
-
|
9 |
-
def extract_frames(video_path):
|
10 |
-
"""
|
11 |
-
Extracts all frames from the input video.
|
12 |
-
"""
|
13 |
-
cap = cv2.VideoCapture(video_path)
|
14 |
-
frames = []
|
15 |
-
while True:
|
16 |
-
ret, frame = cap.read()
|
17 |
-
if not ret:
|
18 |
-
break
|
19 |
-
frames.append(frame)
|
20 |
-
cap.release()
|
21 |
-
print(f"Extracted {len(frames)} frames from {video_path}")
|
22 |
-
return frames
|
23 |
-
|
24 |
-
def apply_style_propagation(frames, style_image_path):
|
25 |
-
"""
|
26 |
-
Applies the style from the provided keyframe image to every frame using optical flow,
|
27 |
-
with additional corrections:
|
28 |
-
- Median filtering of flow components.
|
29 |
-
- Patch-based fallback for blocks with extreme flow.
|
30 |
-
- Temporal reset blending with the original style.
|
31 |
-
- Sharpening after warping.
|
32 |
-
"""
|
33 |
-
# Load and resize the style image to match video dimensions.
|
34 |
-
style_image = cv2.imread(style_image_path)
|
35 |
-
if style_image is None:
|
36 |
-
raise ValueError(f"Failed to load style image from {style_image_path}")
|
37 |
-
h, w = frames[0].shape[:2]
|
38 |
-
style_image = cv2.resize(style_image, (w, h))
|
39 |
-
# Keep a copy for temporal re-anchoring.
|
40 |
-
original_styled = style_image.copy()
|
41 |
-
|
42 |
-
styled_frames = [style_image]
|
43 |
-
prev_gray = cv2.cvtColor(frames[0], cv2.COLOR_BGR2GRAY)
|
44 |
-
|
45 |
-
# Parameters for corrections:
|
46 |
-
reset_interval = 30 # Every 30 frames, blend with original style.
|
47 |
-
block_size = 16 # Size of block for patch matching.
|
48 |
-
patch_threshold = 10 # If mean flow magnitude in a block exceeds this, use patch matching.
|
49 |
-
search_margin = 10 # Margin around block for patch matching.
|
50 |
-
|
51 |
-
for i in tqdm(range(1, len(frames)), desc="Propagating style"):
|
52 |
-
# Compute optical flow between the previous and current grayscale frames.
|
53 |
-
curr_gray = cv2.cvtColor(frames[i], cv2.COLOR_BGR2GRAY)
|
54 |
-
flow = cv2.calcOpticalFlowFarneback(
|
55 |
-
prev_gray, curr_gray, None,
|
56 |
-
pyr_scale=0.5, levels=3, winsize=15,
|
57 |
-
iterations=3, poly_n=5, poly_sigma=1.2, flags=0
|
58 |
-
)
|
59 |
-
|
60 |
-
# --- Method 3: Median filtering of the flow components ---
|
61 |
-
flow_x = flow[..., 0]
|
62 |
-
flow_y = flow[..., 1]
|
63 |
-
flow_x_filtered = cv2.medianBlur(flow_x, 3)
|
64 |
-
flow_y_filtered = cv2.medianBlur(flow_y, 3)
|
65 |
-
flow_filtered = np.dstack((flow_x_filtered, flow_y_filtered))
|
66 |
-
|
67 |
-
# --- Method 4: Patch-based fallback for extreme flow ---
|
68 |
-
flow_corrected = flow_filtered.copy()
|
69 |
-
for by in range(0, h, block_size):
|
70 |
-
for bx in range(0, w, block_size):
|
71 |
-
# Define block region (handle edges)
|
72 |
-
y1, y2 = by, min(by + block_size, h)
|
73 |
-
x1, x2 = bx, min(bx + block_size, w)
|
74 |
-
block_flow = flow_filtered[y1:y2, x1:x2]
|
75 |
-
# Compute mean magnitude in the block.
|
76 |
-
mag = np.sqrt(block_flow[..., 0]**2 + block_flow[..., 1]**2)
|
77 |
-
mean_mag = np.mean(mag)
|
78 |
-
if mean_mag > patch_threshold:
|
79 |
-
# Use patch matching to recalc flow for this block.
|
80 |
-
patch = prev_gray[y1:y2, x1:x2]
|
81 |
-
# Define search region in current frame.
|
82 |
-
sx1 = max(x1 - search_margin, 0)
|
83 |
-
sy1 = max(by - search_margin, 0)
|
84 |
-
sx2 = min(x2 + search_margin, w)
|
85 |
-
sy2 = min(y2 + search_margin, h)
|
86 |
-
search_region = curr_gray[sy1:sy2, sx1:sx2]
|
87 |
-
if search_region.shape[0] < patch.shape[0] or search_region.shape[1] < patch.shape[1]:
|
88 |
-
continue
|
89 |
-
res = cv2.matchTemplate(search_region, patch, cv2.TM_SQDIFF_NORMED)
|
90 |
-
_, _, min_loc, _ = cv2.minMaxLoc(res)
|
91 |
-
best_x = sx1 + min_loc[0]
|
92 |
-
best_y = sy1 + min_loc[1]
|
93 |
-
# Calculate offset relative to block's top-left corner.
|
94 |
-
offset_x = best_x - x1
|
95 |
-
offset_y = best_y - by
|
96 |
-
# Override flow for the entire block.
|
97 |
-
flow_corrected[y1:y2, x1:x2, 0] = offset_x
|
98 |
-
flow_corrected[y1:y2, x1:x2, 1] = offset_y
|
99 |
-
|
100 |
-
# Compute mapping coordinates.
|
101 |
-
grid_x, grid_y = np.meshgrid(np.arange(w), np.arange(h))
|
102 |
-
map_x = grid_x + flow_corrected[..., 0]
|
103 |
-
map_y = grid_y + flow_corrected[..., 1]
|
104 |
-
map_x = np.clip(map_x, 0, w - 1).astype(np.float32)
|
105 |
-
map_y = np.clip(map_y, 0, h - 1).astype(np.float32)
|
106 |
-
|
107 |
-
# Warp the previous styled frame using the computed mapping.
|
108 |
-
warped_styled = cv2.remap(styled_frames[-1], map_x, map_y, interpolation=cv2.INTER_LINEAR)
|
109 |
-
|
110 |
-
# --- Method 2: Temporal Reset/Re-anchoring ---
|
111 |
-
if i % reset_interval == 0:
|
112 |
-
# Blend the current warped result with the original styled keyframe.
|
113 |
-
warped_styled = cv2.addWeighted(warped_styled, 0.7, original_styled, 0.3, 0)
|
114 |
-
|
115 |
-
# --- Method 5: Sharpening Post-Warping ---
|
116 |
-
kernel = np.array([[0, -1, 0],
|
117 |
-
[-1, 5, -1],
|
118 |
-
[0, -1, 0]], dtype=np.float32)
|
119 |
-
warped_sharpened = cv2.filter2D(warped_styled, -1, kernel)
|
120 |
-
|
121 |
-
styled_frames.append(warped_sharpened)
|
122 |
-
prev_gray = curr_gray
|
123 |
-
|
124 |
-
print(f"Propagated style to {len(styled_frames)} frames.")
|
125 |
-
sample_frame = styled_frames[len(styled_frames) // 2]
|
126 |
-
print(f"Sample styled frame mean intensity: {np.mean(sample_frame):.2f}")
|
127 |
-
return styled_frames
|
128 |
-
|
129 |
-
def save_video_cv2(frames, output_path, fps=30):
|
130 |
-
"""
|
131 |
-
Saves a list of frames as a video using OpenCV.
|
132 |
-
"""
|
133 |
-
h, w, _ = frames[0].shape
|
134 |
-
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
135 |
-
writer = cv2.VideoWriter(output_path, fourcc, fps, (w, h))
|
136 |
-
for frame in frames:
|
137 |
-
writer.write(frame)
|
138 |
-
writer.release()
|
139 |
-
size = os.path.getsize(output_path)
|
140 |
-
print(f"Intermediate video saved to {output_path} (size: {size} bytes)")
|
141 |
-
|
142 |
-
def process_video(video_file, style_image_file, fps=30, target_width=0, target_height=0):
|
143 |
-
"""
|
144 |
-
Processes the input video by applying the style image via optical flow propagation,
|
145 |
-
with additional corrections (methods 2, 3, 4, and 5).
|
146 |
-
Optionally downscale the video and style image to the specified resolution.
|
147 |
-
Then re-encodes the video with FFmpeg for web compatibility.
|
148 |
-
|
149 |
-
Parameters:
|
150 |
-
- video_file: The input video file.
|
151 |
-
- style_image_file: The stylized keyframe image.
|
152 |
-
- fps: Output frames per second.
|
153 |
-
- target_width: Target width for downscaling (0 for original).
|
154 |
-
- target_height: Target height for downscaling (0 for original).
|
155 |
-
|
156 |
-
Returns:
|
157 |
-
- Path to the final output video.
|
158 |
-
"""
|
159 |
-
# Get the video file path.
|
160 |
-
video_path = video_file if isinstance(video_file, str) else video_file["name"]
|
161 |
-
|
162 |
-
# Process the style image input.
|
163 |
-
if isinstance(style_image_file, str):
|
164 |
-
style_image_path = style_image_file
|
165 |
-
elif isinstance(style_image_file, dict) and "name" in style_image_file:
|
166 |
-
style_image_path = style_image_file["name"]
|
167 |
-
elif isinstance(style_image_file, np.ndarray):
|
168 |
-
tmp_style = os.path.join(tempfile.gettempdir(), "temp_style_image.jpeg")
|
169 |
-
cv2.imwrite(tmp_style, cv2.cvtColor(style_image_file, cv2.COLOR_RGB2BGR))
|
170 |
-
style_image_path = tmp_style
|
171 |
-
else:
|
172 |
-
return "Error: Unsupported style image format."
|
173 |
-
|
174 |
-
# Extract frames from the video.
|
175 |
-
frames = extract_frames(video_path)
|
176 |
-
if not frames:
|
177 |
-
return "Error: No frames extracted from the video."
|
178 |
-
|
179 |
-
original_h, original_w = frames[0].shape[:2]
|
180 |
-
print(f"Original video resolution: {original_w}x{original_h}")
|
181 |
-
|
182 |
-
# Downscale if target dimensions are provided.
|
183 |
-
if target_width > 0 and target_height > 0:
|
184 |
-
print(f"Downscaling frames to resolution: {target_width}x{target_height}")
|
185 |
-
frames = [cv2.resize(frame, (target_width, target_height)) for frame in frames]
|
186 |
-
else:
|
187 |
-
print("No downscaling applied. Using original resolution.")
|
188 |
-
|
189 |
-
# Propagate the style using our enhanced method.
|
190 |
-
styled_frames = apply_style_propagation(frames, style_image_path)
|
191 |
-
|
192 |
-
# Save intermediate video using OpenCV to a named temporary file.
|
193 |
-
temp_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
|
194 |
-
temp_video_file.close()
|
195 |
-
temp_video_path = temp_video_file.name
|
196 |
-
save_video_cv2(styled_frames, temp_video_path, fps=fps)
|
197 |
-
|
198 |
-
# Re-encode the video using FFmpeg for browser compatibility.
|
199 |
-
output_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
|
200 |
-
output_video_file.close()
|
201 |
-
output_video_path = output_video_file.name
|
202 |
-
|
203 |
-
try:
|
204 |
-
(
|
205 |
-
ffmpeg
|
206 |
-
.input(temp_video_path)
|
207 |
-
.output(output_video_path, vcodec='libx264', pix_fmt='yuv420p', r=fps)
|
208 |
-
.run(overwrite_output=True, quiet=True)
|
209 |
-
)
|
210 |
-
except ffmpeg.Error as e:
|
211 |
-
print("FFmpeg error:", e)
|
212 |
-
return "Error during video re-encoding."
|
213 |
-
|
214 |
-
final_size = os.path.getsize(output_video_path)
|
215 |
-
print(f"Output video saved to {output_video_path} (size: {final_size} bytes)")
|
216 |
-
if final_size == 0:
|
217 |
-
return "Error: Output video file is empty."
|
218 |
-
|
219 |
-
# Clean up the intermediate file.
|
220 |
-
os.remove(temp_video_path)
|
221 |
-
|
222 |
-
return output_video_path
|
223 |
-
|
224 |
-
iface = gr.Interface(
|
225 |
-
fn=process_video,
|
226 |
-
inputs=[
|
227 |
-
gr.Video(label="Input Video (v.mp4)"),
|
228 |
-
gr.Image(label="Stylized Keyframe (a.jpeg)"),
|
229 |
-
gr.Slider(minimum=1, maximum=60, step=1, value=30, label="Output FPS"),
|
230 |
-
gr.Slider(minimum=0, maximum=1920, step=1, value=0, label="Target Width (0 for original)"),
|
231 |
-
gr.Slider(minimum=0, maximum=1080, step=1, value=0, label="Target Height (0 for original)")
|
232 |
-
],
|
233 |
-
outputs=gr.Video(label="Styled Video"),
|
234 |
-
title="Optical Flow Style Propagation with Corrections",
|
235 |
-
description=(
|
236 |
-
"Upload a video and a stylized keyframe image. Optionally downscale both to a target resolution. "
|
237 |
-
"The style is propagated using optical flow with additional corrections:\n"
|
238 |
-
"• Temporal re-anchoring\n"
|
239 |
-
"• Median filtering of the flow\n"
|
240 |
-
"• Patch-based flow correction\n"
|
241 |
-
"• Post-warp sharpening\n"
|
242 |
-
"The output video is re-encoded for web compatibility."
|
243 |
-
)
|
244 |
-
)
|
245 |
-
|
246 |
-
if __name__ == "__main__":
|
247 |
-
iface.launch(share=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|