ghostsInTheMachine
commited on
Commit
•
8df521f
1
Parent(s):
855f2e9
Update app.py
Browse files
app.py
CHANGED
@@ -10,7 +10,7 @@ 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 |
-
import io
|
14 |
|
15 |
# Set up logging
|
16 |
logging.basicConfig(level=logging.INFO)
|
@@ -61,6 +61,8 @@ def process_frame(image, seed=0):
|
|
61 |
# Process video frames and generate depth maps
|
62 |
def process_video(video_path, fps=0, seed=0, batch_size=16):
|
63 |
"""Process video, batch frames, and use L40s GPU to generate depth maps."""
|
|
|
|
|
64 |
try:
|
65 |
start_time = time.time()
|
66 |
|
@@ -76,80 +78,76 @@ def process_video(video_path, fps=0, seed=0, batch_size=16):
|
|
76 |
|
77 |
logger.info(f"Processing {total_frames} frames at {fps} FPS...")
|
78 |
|
79 |
-
# Create
|
80 |
-
|
81 |
-
|
82 |
-
os.makedirs(frames_dir, exist_ok=True)
|
83 |
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
# Process each frame in the batch
|
92 |
-
for frame in frames_batch:
|
93 |
-
depth_map = process_frame(Image.fromarray(frame), seed)
|
94 |
-
depth_maps.append(depth_map)
|
95 |
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
|
|
|
|
|
|
|
|
|
|
102 |
|
103 |
-
|
104 |
-
|
105 |
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
zip_data = f.read()
|
145 |
-
with open(output_video_path, 'rb') as f:
|
146 |
-
video_data = f.read()
|
147 |
-
|
148 |
-
yield None, (zip_filename, zip_data), (video_filename, video_data), f"Processing complete! Total time: {total_time:.2f} seconds"
|
149 |
|
150 |
except Exception as e:
|
151 |
logger.error(f"Error: {e}")
|
152 |
yield None, None, None, f"Error processing video: {e}"
|
|
|
|
|
153 |
|
154 |
# Wrapper function with error handling
|
155 |
def process_wrapper(video, fps=0, seed=0, batch_size=16):
|
@@ -165,7 +163,7 @@ def process_wrapper(video, fps=0, seed=0, batch_size=16):
|
|
165 |
except Exception as e:
|
166 |
raise gr.Error(f"Error processing video: {str(e)}")
|
167 |
|
168 |
-
# Custom CSS for styling
|
169 |
custom_css = """
|
170 |
/* Your existing custom CSS */
|
171 |
"""
|
|
|
10 |
import moviepy.editor as mp
|
11 |
from infer import lotus # Import the depth model inference function
|
12 |
import logging
|
13 |
+
import io
|
14 |
|
15 |
# Set up logging
|
16 |
logging.basicConfig(level=logging.INFO)
|
|
|
61 |
# Process video frames and generate depth maps
|
62 |
def process_video(video_path, fps=0, seed=0, batch_size=16):
|
63 |
"""Process video, batch frames, and use L40s GPU to generate depth maps."""
|
64 |
+
# Create a persistent temporary directory
|
65 |
+
temp_dir = tempfile.mkdtemp()
|
66 |
try:
|
67 |
start_time = time.time()
|
68 |
|
|
|
78 |
|
79 |
logger.info(f"Processing {total_frames} frames at {fps} FPS...")
|
80 |
|
81 |
+
# Create directory for frame sequence and outputs
|
82 |
+
frames_dir = os.path.join(temp_dir, "frames")
|
83 |
+
os.makedirs(frames_dir, exist_ok=True)
|
|
|
84 |
|
85 |
+
processed_frames = []
|
86 |
+
|
87 |
+
# Process frames in batches
|
88 |
+
for i in range(0, total_frames, batch_size):
|
89 |
+
frames_batch = frames[i:i+batch_size]
|
90 |
+
depth_maps = []
|
|
|
|
|
|
|
|
|
|
|
91 |
|
92 |
+
# Process each frame in the batch
|
93 |
+
for frame in frames_batch:
|
94 |
+
depth_map = process_frame(Image.fromarray(frame), seed)
|
95 |
+
depth_maps.append(depth_map)
|
96 |
+
|
97 |
+
for j, depth_map in enumerate(depth_maps):
|
98 |
+
if depth_map is not None:
|
99 |
+
# Save frame
|
100 |
+
frame_index = i + j
|
101 |
+
frame_path = os.path.join(frames_dir, f"frame_{frame_index:06d}.png")
|
102 |
+
Image.fromarray(depth_map).save(frame_path)
|
103 |
|
104 |
+
# Collect processed frame for preview
|
105 |
+
processed_frames.append(depth_map)
|
106 |
|
107 |
+
# Update preview every 10% progress
|
108 |
+
if frame_index % max(1, total_frames // 10) == 0:
|
109 |
+
elapsed_time = time.time() - start_time
|
110 |
+
progress = (frame_index / total_frames) * 100
|
111 |
+
yield processed_frames[-1], None, None, f"Processed {frame_index}/{total_frames} frames... ({progress:.2f}%) Elapsed: {elapsed_time:.2f}s"
|
112 |
+
else:
|
113 |
+
logger.error(f"Error processing frame {frame_index}")
|
114 |
+
|
115 |
+
logger.info("Creating output files...")
|
116 |
|
117 |
+
# Create ZIP of frame sequence
|
118 |
+
zip_filename = f"depth_frames_{int(time.time())}.zip"
|
119 |
+
zip_path = os.path.join(temp_dir, zip_filename)
|
120 |
+
shutil.make_archive(zip_path[:-4], 'zip', frames_dir)
|
121 |
|
122 |
+
# Create MP4 video
|
123 |
+
video_filename = f"depth_video_{int(time.time())}.mp4"
|
124 |
+
output_video_path = os.path.join(temp_dir, video_filename)
|
125 |
|
126 |
+
try:
|
127 |
+
# FFmpeg settings for high-quality MP4
|
128 |
+
(
|
129 |
+
ffmpeg
|
130 |
+
.input(os.path.join(frames_dir, 'frame_%06d.png'), pattern_type='sequence', framerate=fps)
|
131 |
+
.output(output_video_path, vcodec='libx264', pix_fmt='yuv420p', crf=17)
|
132 |
+
.run(overwrite_output=True)
|
133 |
+
)
|
134 |
+
logger.info("MP4 video created successfully!")
|
135 |
|
136 |
+
except ffmpeg.Error as e:
|
137 |
+
logger.error(f"Error creating video: {e.stderr.decode() if e.stderr else str(e)}")
|
138 |
+
output_video_path = None
|
139 |
|
140 |
+
total_time = time.time() - start_time
|
141 |
+
logger.info("Processing complete!")
|
142 |
+
|
143 |
+
# Yield the file paths
|
144 |
+
yield None, zip_path, output_video_path, f"Processing complete! Total time: {total_time:.2f} seconds"
|
|
|
|
|
|
|
|
|
|
|
145 |
|
146 |
except Exception as e:
|
147 |
logger.error(f"Error: {e}")
|
148 |
yield None, None, None, f"Error processing video: {e}"
|
149 |
+
# Do not delete temp_dir here; we need the files to persist
|
150 |
+
# Cleanup can be handled elsewhere if necessary
|
151 |
|
152 |
# Wrapper function with error handling
|
153 |
def process_wrapper(video, fps=0, seed=0, batch_size=16):
|
|
|
163 |
except Exception as e:
|
164 |
raise gr.Error(f"Error processing video: {str(e)}")
|
165 |
|
166 |
+
# Custom CSS for styling (unchanged)
|
167 |
custom_css = """
|
168 |
/* Your existing custom CSS */
|
169 |
"""
|