ghostsInTheMachine
commited on
Commit
•
a85f402
1
Parent(s):
b7ae821
Update app.py
Browse files
app.py
CHANGED
@@ -7,157 +7,151 @@ import time
|
|
7 |
import ffmpeg
|
8 |
import numpy as np
|
9 |
from PIL import Image
|
10 |
-
from concurrent.futures import ThreadPoolExecutor
|
11 |
import moviepy.editor as mp
|
12 |
from infer import lotus # Import the depth model inference function
|
|
|
13 |
|
14 |
-
# Set
|
15 |
-
|
|
|
16 |
|
17 |
-
#
|
|
|
|
|
|
|
18 |
def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)):
|
19 |
-
"""Preprocess the video to resize and
|
20 |
video = mp.VideoFileClip(video_path)
|
21 |
|
22 |
# Resize video if it's larger than the target resolution
|
23 |
if video.size[0] > max_resolution[0] or video.size[1] > max_resolution[1]:
|
24 |
video = video.resize(newsize=max_resolution)
|
25 |
|
26 |
-
#
|
27 |
-
|
|
|
28 |
|
29 |
return video
|
30 |
|
31 |
-
|
|
|
32 |
"""Process a single frame through the depth model and return depth map."""
|
33 |
try:
|
34 |
-
#
|
35 |
-
|
|
|
|
|
|
|
|
|
36 |
|
37 |
-
# Save temporary image (lotus requires a file path)
|
38 |
-
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
|
39 |
-
image.save(tmp.name)
|
40 |
-
|
41 |
-
# Process through the depth model (lotus)
|
42 |
-
_, output_d = lotus(tmp.name, 'depth', seed, device)
|
43 |
-
|
44 |
-
# Clean up temp file
|
45 |
-
os.unlink(tmp.name)
|
46 |
-
|
47 |
# Convert depth output to numpy array
|
48 |
depth_array = np.array(output_d)
|
49 |
return depth_array
|
50 |
|
51 |
except Exception as e:
|
52 |
-
|
53 |
return None
|
54 |
|
55 |
-
|
56 |
-
def process_video(video_path, fps=0, seed=0,
|
57 |
"""Process video, batch frames, and use L40s GPU to generate depth maps."""
|
58 |
-
temp_dir = None
|
59 |
try:
|
60 |
start_time = time.time()
|
61 |
|
62 |
# Preprocess the video
|
63 |
-
video = preprocess_video(video_path)
|
64 |
|
65 |
# Use original video FPS if not specified
|
66 |
if fps == 0:
|
67 |
fps = video.fps
|
68 |
|
69 |
-
frames = list(video.iter_frames(fps=fps))
|
70 |
total_frames = len(frames)
|
71 |
|
72 |
-
|
73 |
-
|
74 |
-
# Create temporary directory for frame sequence
|
75 |
-
temp_dir = tempfile.mkdtemp()
|
76 |
-
frames_dir = os.path.join(temp_dir, "frames")
|
77 |
-
os.makedirs(frames_dir, exist_ok=True)
|
78 |
-
|
79 |
-
# Process frames in larger batches (based on GPU VRAM)
|
80 |
-
batch_size = 50 # Increased batch size to fully utilize the GPU's capabilities
|
81 |
-
processed_frames = []
|
82 |
|
83 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
84 |
for i in range(0, total_frames, batch_size):
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
zip_filename = f"depth_frames_{int(time.time())}.zip"
|
111 |
-
zip_path = os.path.join(output_dir, zip_filename)
|
112 |
-
shutil.make_archive(zip_path[:-4], 'zip', frames_dir)
|
113 |
-
|
114 |
-
# Create MP4 video
|
115 |
-
video_filename = f"depth_video_{int(time.time())}.mp4"
|
116 |
-
video_path = os.path.join(output_dir, video_filename)
|
117 |
-
|
118 |
-
try:
|
119 |
-
# FFmpeg settings for high-quality MP4
|
120 |
-
stream = ffmpeg.input(
|
121 |
-
os.path.join(frames_dir, 'frame_%06d.png'),
|
122 |
-
pattern_type='sequence',
|
123 |
-
framerate=fps
|
124 |
-
)
|
125 |
|
126 |
-
|
127 |
-
stream,
|
128 |
-
video_path,
|
129 |
-
vcodec='libx264',
|
130 |
-
pix_fmt='yuv420p',
|
131 |
-
crf=17, # High quality
|
132 |
-
threads=max_workers
|
133 |
-
)
|
134 |
|
135 |
-
|
136 |
-
|
|
|
|
|
137 |
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
144 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
145 |
except Exception as e:
|
146 |
-
|
147 |
yield None, None, None, f"Error processing video: {e}"
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
shutil.rmtree(temp_dir)
|
152 |
-
except Exception as e:
|
153 |
-
print(f"Error cleaning up temp directory: {e}")
|
154 |
-
|
155 |
-
def process_wrapper(video, fps=0, seed=0, max_workers=32):
|
156 |
if video is None:
|
157 |
raise gr.Error("Please upload a video.")
|
158 |
try:
|
159 |
outputs = []
|
160 |
-
for output in process_video(video, fps, seed,
|
161 |
outputs.append(output)
|
162 |
yield output
|
163 |
return outputs[-1]
|
@@ -205,22 +199,23 @@ with gr.Blocks(css=custom_css) as demo:
|
|
205 |
|
206 |
with gr.Row():
|
207 |
with gr.Column():
|
208 |
-
video_input = gr.Video(label="Upload Video", interactive=True
|
209 |
-
fps_slider = gr.Slider(minimum=0, maximum=60, step=1, value=0, label="Output FPS")
|
210 |
-
seed_slider = gr.
|
211 |
-
|
212 |
-
btn = gr.Button("Process Video"
|
213 |
|
214 |
with gr.Column():
|
215 |
-
preview_image = gr.Image(label="Live Preview"
|
216 |
output_frames_zip = gr.File(label="Download Frame Sequence (ZIP)")
|
217 |
output_video = gr.File(label="Download Video (MP4)")
|
218 |
time_textbox = gr.Textbox(label="Status", interactive=False)
|
219 |
|
220 |
-
btn.click(
|
221 |
-
|
222 |
-
|
223 |
-
|
|
|
224 |
|
225 |
demo.queue()
|
226 |
|
|
|
7 |
import ffmpeg
|
8 |
import numpy as np
|
9 |
from PIL import Image
|
|
|
10 |
import moviepy.editor as mp
|
11 |
from infer import lotus # Import the depth model inference function
|
12 |
+
import logging
|
13 |
|
14 |
+
# Set up logging
|
15 |
+
logging.basicConfig(level=logging.INFO)
|
16 |
+
logger = logging.getLogger(__name__)
|
17 |
|
18 |
+
# Set device to use the L40s GPU explicitly
|
19 |
+
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
20 |
+
|
21 |
+
# Preprocess the video to adjust resolution and frame rate
|
22 |
def preprocess_video(video_path, target_fps=24, max_resolution=(1920, 1080)):
|
23 |
+
"""Preprocess the video to resize and adjust its frame rate."""
|
24 |
video = mp.VideoFileClip(video_path)
|
25 |
|
26 |
# Resize video if it's larger than the target resolution
|
27 |
if video.size[0] > max_resolution[0] or video.size[1] > max_resolution[1]:
|
28 |
video = video.resize(newsize=max_resolution)
|
29 |
|
30 |
+
# Adjust FPS if target_fps is specified
|
31 |
+
if target_fps > 0:
|
32 |
+
video = video.set_fps(target_fps)
|
33 |
|
34 |
return video
|
35 |
|
36 |
+
# Process a single frame through the depth model
|
37 |
+
def process_frame(image, seed=0):
|
38 |
"""Process a single frame through the depth model and return depth map."""
|
39 |
try:
|
40 |
+
# Set seeds for reproducibility
|
41 |
+
torch.manual_seed(seed)
|
42 |
+
np.random.seed(seed)
|
43 |
+
|
44 |
+
# Process through the depth model (assuming lotus accepts image data)
|
45 |
+
_, output_d = lotus(image, 'depth', seed, device)
|
46 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
47 |
# Convert depth output to numpy array
|
48 |
depth_array = np.array(output_d)
|
49 |
return depth_array
|
50 |
|
51 |
except Exception as e:
|
52 |
+
logger.error(f"Error processing frame: {e}")
|
53 |
return None
|
54 |
|
55 |
+
# Process video frames and generate depth maps
|
56 |
+
def process_video(video_path, fps=0, seed=0, batch_size=16):
|
57 |
"""Process video, batch frames, and use L40s GPU to generate depth maps."""
|
|
|
58 |
try:
|
59 |
start_time = time.time()
|
60 |
|
61 |
# Preprocess the video
|
62 |
+
video = preprocess_video(video_path, target_fps=fps)
|
63 |
|
64 |
# Use original video FPS if not specified
|
65 |
if fps == 0:
|
66 |
fps = video.fps
|
67 |
|
68 |
+
frames = list(video.iter_frames(fps=video.fps))
|
69 |
total_frames = len(frames)
|
70 |
|
71 |
+
logger.info(f"Processing {total_frames} frames at {fps} FPS...")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
72 |
|
73 |
+
# Create temporary directory for frame sequence and outputs
|
74 |
+
with tempfile.TemporaryDirectory() as temp_dir:
|
75 |
+
frames_dir = os.path.join(temp_dir, "frames")
|
76 |
+
os.makedirs(frames_dir, exist_ok=True)
|
77 |
+
|
78 |
+
processed_frames = []
|
79 |
+
|
80 |
+
# Process frames in batches
|
81 |
for i in range(0, total_frames, batch_size):
|
82 |
+
frames_batch = frames[i:i+batch_size]
|
83 |
+
depth_maps = []
|
84 |
+
|
85 |
+
# Process each frame in the batch
|
86 |
+
for frame in frames_batch:
|
87 |
+
depth_map = process_frame(Image.fromarray(frame), seed)
|
88 |
+
depth_maps.append(depth_map)
|
89 |
+
|
90 |
+
for j, depth_map in enumerate(depth_maps):
|
91 |
+
if depth_map is not None:
|
92 |
+
# Save frame
|
93 |
+
frame_index = i + j
|
94 |
+
frame_path = os.path.join(frames_dir, f"frame_{frame_index:06d}.png")
|
95 |
+
Image.fromarray(depth_map).save(frame_path)
|
96 |
+
|
97 |
+
# Collect processed frame for preview
|
98 |
+
processed_frames.append(depth_map)
|
99 |
+
|
100 |
+
# Update preview every 10% progress
|
101 |
+
if frame_index % max(1, total_frames // 10) == 0:
|
102 |
+
elapsed_time = time.time() - start_time
|
103 |
+
progress = (frame_index / total_frames) * 100
|
104 |
+
yield processed_frames[-1], None, None, f"Processed {frame_index}/{total_frames} frames... ({progress:.2f}%) Elapsed: {elapsed_time:.2f}s"
|
105 |
+
else:
|
106 |
+
logger.error(f"Error processing frame {frame_index}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
107 |
|
108 |
+
logger.info("Creating output files...")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
109 |
|
110 |
+
# Create ZIP of frame sequence
|
111 |
+
zip_filename = f"depth_frames_{int(time.time())}.zip"
|
112 |
+
zip_path = os.path.join(temp_dir, zip_filename)
|
113 |
+
shutil.make_archive(zip_path[:-4], 'zip', frames_dir)
|
114 |
|
115 |
+
# Create MP4 video
|
116 |
+
video_filename = f"depth_video_{int(time.time())}.mp4"
|
117 |
+
output_video_path = os.path.join(temp_dir, video_filename)
|
118 |
+
|
119 |
+
try:
|
120 |
+
# FFmpeg settings for high-quality MP4
|
121 |
+
(
|
122 |
+
ffmpeg
|
123 |
+
.input(os.path.join(frames_dir, 'frame_%06d.png'), pattern_type='sequence', framerate=fps)
|
124 |
+
.output(output_video_path, vcodec='libx264', pix_fmt='yuv420p', crf=17)
|
125 |
+
.run(overwrite_output=True)
|
126 |
+
)
|
127 |
+
logger.info("MP4 video created successfully!")
|
128 |
+
|
129 |
+
except ffmpeg.Error as e:
|
130 |
+
logger.error(f"Error creating video: {e.stderr.decode() if e.stderr else str(e)}")
|
131 |
+
output_video_path = None
|
132 |
|
133 |
+
total_time = time.time() - start_time
|
134 |
+
logger.info("Processing complete!")
|
135 |
+
|
136 |
+
# Read output files to return as bytes
|
137 |
+
with open(zip_path, 'rb') as f:
|
138 |
+
zip_data = f.read()
|
139 |
+
with open(output_video_path, 'rb') as f:
|
140 |
+
video_data = f.read()
|
141 |
+
|
142 |
+
yield None, (zip_filename, zip_data), (video_filename, video_data), f"Processing complete! Total time: {total_time:.2f} seconds"
|
143 |
+
|
144 |
except Exception as e:
|
145 |
+
logger.error(f"Error: {e}")
|
146 |
yield None, None, None, f"Error processing video: {e}"
|
147 |
+
|
148 |
+
# Wrapper function with error handling
|
149 |
+
def process_wrapper(video, fps=0, seed=0, batch_size=16):
|
|
|
|
|
|
|
|
|
|
|
150 |
if video is None:
|
151 |
raise gr.Error("Please upload a video.")
|
152 |
try:
|
153 |
outputs = []
|
154 |
+
for output in process_video(video.name, fps, seed, batch_size):
|
155 |
outputs.append(output)
|
156 |
yield output
|
157 |
return outputs[-1]
|
|
|
199 |
|
200 |
with gr.Row():
|
201 |
with gr.Column():
|
202 |
+
video_input = gr.Video(label="Upload Video", interactive=True)
|
203 |
+
fps_slider = gr.Slider(minimum=0, maximum=60, step=1, value=0, label="Output FPS (0 for original)")
|
204 |
+
seed_slider = gr.Number(value=0, label="Seed")
|
205 |
+
batch_size_slider = gr.Slider(minimum=1, maximum=64, step=1, value=16, label="Batch Size")
|
206 |
+
btn = gr.Button("Process Video")
|
207 |
|
208 |
with gr.Column():
|
209 |
+
preview_image = gr.Image(label="Live Preview")
|
210 |
output_frames_zip = gr.File(label="Download Frame Sequence (ZIP)")
|
211 |
output_video = gr.File(label="Download Video (MP4)")
|
212 |
time_textbox = gr.Textbox(label="Status", interactive=False)
|
213 |
|
214 |
+
btn.click(
|
215 |
+
fn=process_wrapper,
|
216 |
+
inputs=[video_input, fps_slider, seed_slider, batch_size_slider],
|
217 |
+
outputs=[preview_image, output_frames_zip, output_video, time_textbox]
|
218 |
+
)
|
219 |
|
220 |
demo.queue()
|
221 |
|