refactored and added schedule

app.py

@@ -11,6 +11,10 @@ import numpy as np
 import os
 import tempfile
 import uuid
+import schedule
+import time
+import threading
+import shutil
 
 torch.set_float32_matmul_precision("medium")
 
@@ -29,64 +33,116 @@ transform_image = transforms.Compose(
 )
 
 
+def load_background_image(bg, image_size):
+ """Loads and resizes the background image based on the provided input."""
+ try:
+ if isinstance(bg, str) and bg.startswith("#"):
+ color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
+ return Image.new("RGBA", image_size, color_rgb + (255,))
+ elif isinstance(bg, Image.Image):
+ return bg.convert("RGBA").resize(image_size)
+ else:
+ return Image.open(bg).convert("RGBA").resize(image_size)
+ except Exception as e:
+ print(f"Error opening background image: {e}")
+ return None
+
+def clear_temp_directory():
+ temp_dir = "temp"
+ if os.path.exists(temp_dir):
+ for filename in os.listdir(temp_dir):
+ file_path = os.path.join(temp_dir, filename)
+ try:
+ if os.path.isfile(file_path) or os.path.islink(file_path):
+ os.unlink(file_path)
+ elif os.path.isdir(file_path):
+ shutil.rmtree(file_path)
+ except Exception as e:
+ print(f"Failed to delete {file_path}. Reason: {e}")
+
+def run_scheduler():
+ schedule.every(10).minutes.do(clear_temp_directory)
+ while True:
+ schedule.run_pending()
+ time.sleep(1)
+
+# Start the scheduler in a separate thread
+scheduler_thread = threading.Thread(target=run_scheduler)
+scheduler_thread.daemon = True # Allow the main thread to exit even if the scheduler is running
+scheduler_thread.start()
+
 @spaces.GPU
-def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
+def process_video(input_video_path, bg_type="Color", bg_image_path=None, bg_video_path=None, 
+ bg_color="#00FF00", output_fps=0, video_handling_mode="slow_down"):
+ """Processes the input video and replaces the background."""
  try:
  # Load the video using moviepy
- video = mp.VideoFileClip(vid)
+ video = mp.VideoFileClip(input_video_path)
 
  # Load original fps if fps value is equal to 0
- if fps == 0:
- fps = video.fps
+ if output_fps == 0:
+ output_fps = video.fps
 
  # Extract audio from the video
  audio = video.audio
 
  # Extract frames at the specified FPS
- frames = video.iter_frames(fps=fps)
+ frames = video.iter_frames(fps=output_fps)
 
  # Process each frame for background removal
  processed_frames = []
- yield gr.update(visible=True), gr.update(visible=False)
+ yield gr.update(visible=True), gr.update(visible=False), gr.update(value=0)
 
  if bg_type == "Video":
- background_video = mp.VideoFileClip(bg_video)
+ background_video = mp.VideoFileClip(bg_video_path)
  if background_video.duration < video.duration:
- if video_handling == "slow_down":
+ if video_handling_mode == "slow_down":
  background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
- else: # video_handling == "loop"
+ else: # video_handling_mode == "loop"
  background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
- background_frames = list(background_video.iter_frames(fps=fps)) # Convert to list
+ background_frames = list(background_video.iter_frames(fps=output_fps)) # Convert to list
  else:
  background_frames = None
 
  bg_frame_index = 0 # Initialize background frame index
-
- for i, frame in enumerate(frames):
- pil_image = Image.fromarray(frame)
- if bg_type == "Color":
- processed_image = process(pil_image, color)
- elif bg_type == "Image":
- processed_image = process(pil_image, bg_image)
- elif bg_type == "Video":
- if video_handling == "slow_down":
- background_frame = background_frames[bg_frame_index % len(background_frames)]
- bg_frame_index += 1
- background_image = Image.fromarray(background_frame)
+ total_frames = len(list(frames))
+ frames = video.iter_frames(fps=output_fps)
+ batch_size = 4 
+
+ for i in range(0, total_frames, batch_size):
+ batch_frames = list(frames)[i:i+batch_size]
+ processed_batch = []
+ for frame in batch_frames:
+ pil_image = Image.fromarray(frame)
+ if bg_type == "Color":
+ background_image = load_background_image(bg_color, pil_image.size)
+ elif bg_type == "Image":
+ background_image = load_background_image(bg_image_path, pil_image.size)
+ elif bg_type == "Video":
+ if video_handling_mode == "slow_down":
+ background_frame = background_frames[bg_frame_index % len(background_frames)]
+ bg_frame_index += 1
+ background_image = Image.fromarray(background_frame)
+ else: # video_handling_mode == "loop"
+ background_frame = background_frames[bg_frame_index % len(background_frames)]
+ bg_frame_index += 1
+ background_image = Image.fromarray(background_frame)
+ else:
+ background_image = None # Default to original image if no background is selected
+
+ if background_image is not None:
  processed_image = process(pil_image, background_image)
- else: # video_handling == "loop"
- background_frame = background_frames[bg_frame_index % len(background_frames)]
- bg_frame_index += 1
- background_image = Image.fromarray(background_frame)
- processed_image = process(pil_image, background_image)
- else:
- processed_image = pil_image # Default to original image if no background is selected
+ else:
+ processed_image = pil_image
+
+ processed_batch.append(np.array(processed_image))
 
- processed_frames.append(np.array(processed_image))
- yield processed_image, None
+ processed_frames.extend(processed_batch)
+ progress = (i + len(batch_frames)) / total_frames
+ yield processed_batch[-1], None, gr.update(value=progress) # Update progress bar
 
  # Create a new video from the processed frames
- processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
+ processed_video = mp.ImageSequenceClip(processed_frames, fps=output_fps)
 
  # Add the original audio back to the processed video
  processed_video = processed_video.set_audio(audio)
@@ -98,15 +154,15 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
  temp_filepath = os.path.join(temp_dir, unique_filename)
  processed_video.write_videofile(temp_filepath, codec="libx264")
 
- yield gr.update(visible=False), gr.update(visible=True)
+ yield gr.update(visible=False), gr.update(visible=True), gr.update(value=1)
  # Return the path to the temporary file
- yield processed_image, temp_filepath
+ yield processed_batch[-1], temp_filepath
 
  except Exception as e:
- print(f"Error: {e}")
- yield gr.update(visible=False), gr.update(visible=True)
- yield None, f"Error processing video: {e}"
-
+ error_message = f"Error processing video: {e}"
+ print(error_message)
+ yield gr.update(visible=False), gr.update(visible=True), gr.update(value=0)
+ yield None, error_message
 
 
 def process(image, bg):
@@ -119,16 +175,8 @@ def process(image, bg):
  pred_pil = transforms.ToPILImage()(pred)
  mask = pred_pil.resize(image_size)
 
- if isinstance(bg, str) and bg.startswith("#"):
- color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
- background = Image.new("RGBA", image_size, color_rgb + (255,))
- elif isinstance(bg, Image.Image):
- background = bg.convert("RGBA").resize(image_size)
- else:
- background = Image.open(bg).convert("RGBA").resize(image_size)
-
  # Composite the image onto the background using the mask
- image = Image.composite(image, background, mask)
+ image = Image.composite(image, bg, mask)
 
  return image
 
@@ -140,6 +188,7 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
  stream_image = gr.Image(label="Streaming Output", visible=False)
  out_video = gr.Video(label="Final Output Video")
  submit_button = gr.Button("Change Background", interactive=True)
+ progress_bar = gr.ProgressBar(label="Processing Progress")
  with gr.Row():
  fps_slider = gr.Slider(
  minimum=0,
@@ -177,17 +226,17 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
  ["rickroll-2sec.mp4", "Color", None, None],
  ],
  inputs=[in_video, bg_type, bg_image, bg_video],
- outputs=[stream_image, out_video],
- fn=fn,
+ outputs=[stream_image, out_video, progress_bar],
+ fn=process_video,
  cache_examples=True,
  cache_mode="eager",
  )
 
 
  submit_button.click(
- fn,
+ process_video,
  inputs=[in_video, bg_type, bg_image, bg_video, color_picker, fps_slider, video_handling_radio],
- outputs=[stream_image, out_video],
+ outputs=[stream_image, out_video, progress_bar],
  )
 
 if __name__ == "__main__":