app.py

import gradio as gr
from loadimg import load_img
import spaces
from transformers import AutoModelForImageSegmentation
import torch
from torchvision import transforms
import moviepy.editor as mp
from PIL import Image
import numpy as np
import os
import tempfile
import uuid
from concurrent.futures import ThreadPoolExecutor

torch.set_float32_matmul_precision("highest")

birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
).to("cuda")

transform_image = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

BATCH_SIZE = 3
executor = ThreadPoolExecutor(max_workers=4)  # Adjust as needed

def get_background_image(bg_type, bg_image, background_frames, current_frame_index, video_handling, slow_down_factor):
    if bg_type == "Video":
        if video_handling == "slow_down":
            frame_index = int(current_frame_index / slow_down_factor)
        else:
            frame_index = current_frame_index
        return Image.fromarray(background_frames[frame_index % len(background_frames)])
    elif bg_type == "Image":
        return bg_image  # Directly returns the image path
    else:  # bg_type == "Color"
        return bg_image  # bg_image here is the color string

@spaces.GPU
def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
    try:
        video = mp.VideoFileClip(vid)
        try:
            audio = video.audio
        except AttributeError:
            audio = None

        if fps == 0:
            fps = video.fps
        frames = video.iter_frames(fps=fps)
        processed_frames = []
        yield gr.update(visible=True), gr.update(visible=False)  # Update Gradio display

        if bg_type == "Video":
            background_video = mp.VideoFileClip(bg_video)
            if background_video.duration < video.duration and video_handling == "slow_down":
                slow_down_factor = video.duration / background_video.duration
            else:
                slow_down_factor = 1
            background_frames = list(background_video.iter_frames(fps=fps))
        else:
            background_frames = None
            slow_down_factor = None  # Not needed for image or color backgrounds

        bg_frame_index = 0
        frame_batch = []

        for i, frame in enumerate(frames):
            frame_batch.append(frame)
            if len(frame_batch) == BATCH_SIZE or i == int(video.fps * video.duration) - 1:
                pil_images = [Image.fromarray(f) for f in frame_batch]

                if bg_type == "Video":
                    processed_images = list(executor.map(process, pil_images, [get_background_image(bg_type, bg_image, background_frames, bg_frame_index + j, video_handling, slow_down_factor) for j in range(len(pil_images))]))
                    bg_frame_index += len(frame_batch)
                elif bg_type == "Color":
                    processed_images = list(executor.map(process, pil_images, [color] * len(pil_images)))  # Use color directly
                elif bg_type == "Image":
                    processed_images = list(executor.map(process, pil_images, [bg_image] * len(pil_images))) # Use image path directly
                else:
                    processed_images = pil_images  # No processing needed

                for processed_image in processed_images:
                    processed_frames.append(np.array(processed_image))
                    yield processed_image, None  # Update Gradio with processed images
                frame_batch = []


        processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
        if audio:
            processed_video = processed_video.set_audio(audio)

        # Save processed video to a temporary file
        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", logger=None)


        yield gr.update(visible=False), gr.update(visible=True) # Update Gradio display
        yield processed_image, temp_filepath # Return final output

    except Exception as e:
        print(f"Error: {e}")
        yield gr.update(visible=False), gr.update(visible=True)
        yield None, f"Error processing video: {e}"




def process(image, bg):
    image_size = image.size
    input_images = transform_image(image).unsqueeze(0).to("cuda")
    with torch.no_grad():
        preds = birefnet(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(image_size)

    if isinstance(bg, str) and bg.startswith("#"):  # If bg is a color
        color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
        background = Image.new("RGBA", image_size, color_rgb + (255,)) # Create image with color
    elif isinstance(bg, Image.Image):
        background = bg.convert("RGBA").resize(image_size) #Resize if bg is an image
    else: #If bg is an image path
        background = Image.open(bg).convert("RGBA").resize(image_size) # Open and resize image

    image = Image.composite(image, background, mask)
    return image



with gr.Blocks(theme=gr.themes.Ocean()) as demo:
    with gr.Row():
        in_video = gr.Video(label="Input Video", interactive=True)
        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)
    with gr.Row():
        fps_slider = gr.Slider(
            minimum=0,
            maximum=60,
            step=1,
            value=0,
            label="Output FPS (0 will inherit the original fps value)",
            interactive=True
        )
        bg_type = gr.Radio(["Color", "Image", "Video"], label="Background Type", value="Color", interactive=True)
        color_picker = gr.ColorPicker(label="Background Color", value="#00FF00", visible=True, interactive=True)
        bg_image = gr.Image(label="Background Image", type="filepath", visible=False, interactive=True)
        bg_video = gr.Video(label="Background Video", visible=False, interactive=True)
        with gr.Column(visible=False) as video_handling_options:
            video_handling_radio = gr.Radio(["slow_down", "loop"], label="Video Handling", value="slow_down", interactive=True)

    def update_visibility(bg_type):
        if bg_type == "Color":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        elif bg_type == "Image":
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
        elif bg_type == "Video":
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
        else:
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)


    bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image, bg_video, video_handling_options])


    examples = gr.Examples(
        [
            ["rickroll-2sec.mp4", "Video", None, "background.mp4"],
            ["rickroll-2sec.mp4", "Image", "images.webp", None],
            ["rickroll-2sec.mp4", "Color", None, None],
        ],
        inputs=[in_video, bg_type, bg_image, bg_video],
        outputs=[stream_image, out_video],
        fn=fn,
        cache_examples=True,
        cache_mode="eager",
    )


    submit_button.click(
        fn,
        inputs=[in_video, bg_type, bg_image, bg_video, color_picker, fps_slider, video_handling_radio],
        outputs=[stream_image, out_video],
    )

if __name__ == "__main__":
    demo.launch(show_error=True)