Update app.py

app.py

@@ -13,7 +13,6 @@ import tempfile
 import uuid
 import schedule
 import time
-import threading
 import shutil
 
 torch.set_float32_matmul_precision("medium")
@@ -33,116 +32,64 @@ 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 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."""
+def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
     try:
         # Load the video using moviepy
-        video = mp.VideoFileClip(input_video_path)
+        video = mp.VideoFileClip(vid)
 
         # Load original fps if fps value is equal to 0
-        if output_fps == 0:
-            output_fps = video.fps
+        if fps == 0:
+            fps = video.fps
 
         # Extract audio from the video
         audio = video.audio
 
         # Extract frames at the specified FPS
-        frames = video.iter_frames(fps=output_fps)
+        frames = video.iter_frames(fps=fps)
 
         # Process each frame for background removal
         processed_frames = []
-        yield gr.update(visible=True), gr.update(visible=False), gr.update(value=0)
+        yield gr.update(visible=True), gr.update(visible=False)
 
         if bg_type == "Video":
-            background_video = mp.VideoFileClip(bg_video_path)
+            background_video = mp.VideoFileClip(bg_video)
             if background_video.duration < video.duration:
-                if video_handling_mode == "slow_down":
+                if video_handling == "slow_down":
                     background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
-                else:  # video_handling_mode == "loop"
+                else:  # video_handling == "loop"
                     background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
-            background_frames = list(background_video.iter_frames(fps=output_fps))  # Convert to list
+            background_frames = list(background_video.iter_frames(fps=fps))  # Convert to list
         else:
             background_frames = None
 
         bg_frame_index = 0  # Initialize background frame index
-        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:
-                    processed_image = pil_image
 
-                processed_batch.append(np.array(processed_image))
+        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)
+                    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
 
-            processed_frames.extend(processed_batch)
-            progress = (i + len(batch_frames)) / total_frames
-            yield processed_batch[-1], None, gr.update(value=progress)  # Update progress bar
+            processed_frames.append(np.array(processed_image))
+            yield processed_image, None
 
         # Create a new video from the processed frames
-        processed_video = mp.ImageSequenceClip(processed_frames, fps=output_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)
@@ -154,15 +101,15 @@ def process_video(input_video_path, bg_type="Color", bg_image_path=None, bg_vide
         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), gr.update(value=1)
+        yield gr.update(visible=False), gr.update(visible=True)
         # Return the path to the temporary file
-        yield processed_batch[-1], temp_filepath
+        yield processed_image, temp_filepath
 
     except Exception as 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
+        print(f"Error: {e}")
+        yield gr.update(visible=False), gr.update(visible=True)
+        yield None, f"Error processing video: {e}"
+
 
 
 def process(image, bg):
@@ -175,11 +122,31 @@ 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, bg, mask)
+    image = Image.composite(image, background, mask)
 
     return image
 
+def clear_temp_directory():
+    temp_dir = "temp"
+    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('Failed to delete %s. Reason: %s' % (file_path, e)) # Keep this print statement for debugging purposes
+
 
 with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     gr.Markdown("# Video Background Remover & Changer\n### You can replace image background with any color, image or video.\nNOTE: As this Space is running on ZERO GPU it has limit. It can handle approx 200frmaes at once. So, if you have big video than use small chunks or Duplicate this space.")
@@ -188,7 +155,6 @@ 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,
@@ -226,18 +192,24 @@ 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, progress_bar],
-        fn=process_video,
+        outputs=[stream_image, out_video],
+        fn=fn,
         cache_examples=True,
         cache_mode="eager",
     )
 
 
     submit_button.click(
-        process_video,
+        fn,
         inputs=[in_video, bg_type, bg_image, bg_video, color_picker, fps_slider, video_handling_radio],
-        outputs=[stream_image, out_video, progress_bar],
+        outputs=[stream_image, out_video],
     )
 
 if __name__ == "__main__":
+    schedule.every(10).minutes.do(clear_temp_directory)
+
+    while True:
+        schedule.run_pending()
+        time.sleep(1)
+
     demo.launch(show_error=True)