Update app.py

app.py

@@ -26,6 +26,7 @@ transform_image = transforms.Compose(
  ]
 )
 
+
 @spaces.GPU
 def fn(vid, bg_type="Color", bg_image=None, color="#00FF00", fps=0):
  # Load the video using moviepy
@@ -38,68 +39,37 @@ def fn(vid, bg_type="Color", bg_image=None, color="#00FF00", fps=0):
  # Extract audio from the video
  audio = video.audio
 
- # Process video in chunks of 1 second
- chunk_duration = 1 # seconds
- total_duration = video.duration
- start_time = 0
- 
- progress = f'<div class="progress-container"><div class="progress-bar" style="--current: {start_time}; --total: {total_duration};"></div></div>'
+ # Extract frames at the specified FPS
+ frames = video.iter_frames(fps=fps)
 
+ # Process each frame for background removal
  processed_frames = []
- yield gr.update(visible=True), gr.update(visible=False), progress
-
- while start_time < total_duration:
- end_time = min(start_time + chunk_duration, total_duration)
- chunk = video.subclip(start_time, end_time)
- chunk_frames = chunk.iter_frames(fps=fps)
-
- for frame in chunk_frames:
- pil_image = Image.fromarray(frame)
- if bg_type == "Color":
- processed_image = process(pil_image, color)
- else:
- processed_image = process(pil_image, bg_image)
- processed_frames.append(np.array(processed_image))
- yield processed_image, None, progress
-
- # Save processed frames for the current chunk
- temp_dir = "temp"
- os.makedirs(temp_dir, exist_ok=True)
- for i, frame in enumerate(processed_frames):
- Image.fromarray(frame).save(os.path.join(temp_dir, f"frame_{start_time}_{i}.png"))
-
- # Clear processed frames for the current chunk
- processed_frames = []
- progress = f'<div class="progress-container"><div class="progress-bar" style="--current: {start_time}; --total: {total_duration};"></div></div>'
-
- yield None, None, progress
-
- start_time += chunk_duration
-
- # Load all saved frames
- all_frames = []
- for filename in sorted(os.listdir(temp_dir)):
- if filename.startswith("frame_") and filename.endswith(".png"):
- frame = np.array(Image.open(os.path.join(temp_dir, filename)))
- all_frames.append(frame)
+ yield gr.update(visible=True), gr.update(visible=False)
+ for frame in frames:
+ pil_image = Image.fromarray(frame)
+ if bg_type == "Color":
+ processed_image = process(pil_image, color)
+ else:
+ processed_image = process(pil_image, bg_image)
+ processed_frames.append(np.array(processed_image))
+ yield processed_image, None
 
  # Create a new video from the processed frames
- processed_video = mp.ImageSequenceClip(all_frames, fps=fps)
+ processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
 
  # Add the original audio back to the processed video
  processed_video = processed_video.set_audio(audio)
 
  # Save the processed video to a temporary file
- temp_filepath = os.path.join(temp_dir, "processed_video.mp4")
+ temp_dir = "temp"
+ os.makedirs(temp_dir, exist_ok=True)
+ unique_filename = str(uuid.uuid4()) + ".mp4"
+ temp_filepath = os.path.join(temp_dir, unique_filename)
  processed_video.write_videofile(temp_filepath, codec="libx264")
 
- # Clean up temporary files
- for filename in os.listdir(temp_dir):
- os.remove(os.path.join(temp_dir, filename))
-
- yield gr.update(visible=False), gr.update(visible=True), progress
+ yield gr.update(visible=False), gr.update(visible=True)
  # Return the path to the temporary file
- yield processed_image, temp_filepath, progress
+ yield processed_image, temp_filepath
 
 
 def process(image, bg):
@@ -124,12 +94,7 @@ def process(image, bg):
  return image
 
 
-css="""
-.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
-.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
-"""
-
-with gr.Blocks(css=css, theme="ocean") as demo:
+with gr.Blocks(theme=gr.themes.Ocean(primary_hue="blue")) as demo:
  with gr.Row():
  in_video = gr.Video(label="Input Video")
  stream_image = gr.Image(label="Streaming Output", visible=False)
@@ -155,12 +120,11 @@ with gr.Blocks(css=css, theme="ocean") as demo:
 
  bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image])
 
- progress_bar = gr.Markdown(elem_id="progress")
 
  examples = gr.Examples(
  [["rickroll-2sec.mp4", "Image", "images.webp"], ["rickroll-2sec.mp4", "Color", None]],
  inputs=[in_video, bg_type, bg_image],
- outputs=[stream_image, out_video, progress_bar],
+ outputs=[stream_image, out_video],
  fn=fn,
  cache_examples=True,
  cache_mode="eager",
@@ -170,7 +134,7 @@ with gr.Blocks(css=css, theme="ocean") as demo:
  submit_button.click(
  fn,
  inputs=[in_video, bg_type, bg_image, color_picker, fps_slider],
- outputs=[stream_image, out_video, progress_bar],
+ outputs=[stream_image, out_video],
  )
 
 if __name__ == "__main__":