app.py

import gradio as gr
import numpy as np
import subprocess
import supervision as sv
import torch
import uuid
from PIL import Image
from tqdm import tqdm
from transformers import pipeline, CLIPModel, CLIPProcessor
from typing import Tuple, List

MARKDOWN = """
# Auto ⚡ ProPainter 🧑‍🎨
This is a demo for automatic removal of objects from videos using
[Segment Anything Model](https://github.com/facebookresearch/segment-anything),
[MetaCLIP](https://github.com/facebookresearch/MetaCLIP), and 
[ProPainter](https://github.com/sczhou/ProPainter) combo.

- [x] Automated object masking using SAM + MetaCLIP
- [x] Automated inpainting using ProPainter
- [ ] Automated ⚡ object masking using FastSAM + MetaCLIP
"""
EXAMPLES = [
    ["https://media.roboflow.com/supervision/video-examples/ball-juggling.mp4", "person", 0.6]
]

START_FRAME = 0
END_FRAME = 10
TOTAL = END_FRAME - START_FRAME
MINIMUM_AREA = 0.01

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SAM_GENERATOR = pipeline(
    task="mask-generation",
    model="facebook/sam-vit-large",
    device=DEVICE)
CLIP_MODEL = CLIPModel.from_pretrained("facebook/metaclip-b32-400m").to(DEVICE)
CLIP_PROCESSOR = CLIPProcessor.from_pretrained("facebook/metaclip-b32-400m")


def run_sam(frame: np.ndarray) -> sv.Detections:
    # convert from Numpy BGR to PIL RGB
    image = Image.fromarray(frame[:, :, ::-1])
    outputs = SAM_GENERATOR(image)
    mask = np.array(outputs['masks'])
    return sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)


def run_clip(frame: np.ndarray, text: List[str]) -> np.ndarray:
    # convert from Numpy BGR to PIL RGB
    image = Image.fromarray(frame[:, :, ::-1])
    inputs = CLIP_PROCESSOR(text=text, images=image, return_tensors="pt").to(DEVICE)
    outputs = CLIP_MODEL(**inputs)
    probs = outputs.logits_per_image.softmax(dim=1)
    return probs.detach().cpu().numpy()


def gray_background(image: np.ndarray, mask: np.ndarray, gray_value=128):
    gray_color = np.array([gray_value, gray_value, gray_value], dtype=np.uint8)
    return np.where(mask[..., None], image, gray_color)


def filter_detections_by_area(frame: np.ndarray, detections: sv.Detections, minimum_area: float) -> sv.Detections:
    frame_width, frame_height = frame.shape[1], frame.shape[0]
    frame_area = frame_width * frame_height
    return detections[detections.area > minimum_area * frame_area]


def filter_detections_by_prompt(frame: np.ndarray, detections: sv.Detections, prompt: str, confidence: float) -> sv.Detections:
    text = [f"a picture of {prompt}", "a picture of background"]
    filtering_mask = []
    for xyxy, mask in zip(detections.xyxy, detections.mask):
        crop = gray_background(
            image=sv.crop_image(image=frame, xyxy=xyxy),
            mask=sv.crop_image(image=mask, xyxy=xyxy))
        probs = run_clip(frame=crop, text=text)
        filtering_mask.append(probs[0][0] > confidence)

    return detections[np.array(filtering_mask)]


def mask_frame(frame: np.ndarray, prompt: str, confidence: float) -> np.ndarray:
    detections = run_sam(frame)
    detections = filter_detections_by_area(
        frame=frame, detections=detections, minimum_area=MINIMUM_AREA)
    detections = filter_detections_by_prompt(
        frame=frame, detections=detections, prompt=prompt, confidence=confidence)
    # converting set of masks to a single mask
    mask = np.any(detections.mask, axis=0).astype(np.uint8) * 255
    # converting single channel mask to 3 channel mask
    return np.repeat(mask[:, :, np.newaxis], 3, axis=2)


def mask_video(source_video: str, prompt: str, confidence: float, frames_dir: str, masked_frames_dir: str) -> None:
    frame_iterator = iter(sv.get_video_frames_generator(
        source_path=source_video, start=START_FRAME, end=END_FRAME))

    with sv.ImageSink(masked_frames_dir, image_name_pattern="{:05d}.png") as masked_frames_sink:
        with sv.ImageSink(frames_dir, image_name_pattern="{:05d}.jpg") as frames_sink:
            for _ in tqdm(range(TOTAL), desc="Masking frames"):
                frame = next(frame_iterator)
                frames_sink.save_image(frame)
                masked_frame = mask_frame(frame, prompt, confidence)
                masked_frames_sink.save_image(masked_frame)

    return frames_dir, masked_frames_dir


def execute_command(command: str) -> None:
    subprocess.run(command, check=True)


def paint_video(frames_dir: str, masked_frames_dir: str, results_dir: str) -> None:
    command = [
        f"python",
        f"inference_propainter.py",
        f"--video={frames_dir}",
        f"--mask={masked_frames_dir}",
        f"--output={results_dir}",
        f"--save_fps={25}"
    ]
    execute_command(command)


def process(
    source_video: str,
    prompt: str,
    confidence: float,
    progress=gr.Progress(track_tqdm=True)
) -> Tuple[str, str]:
    name = str(uuid.uuid4())
    frames_dir = f"{name}/frames"
    masked_frames_dir = f"{name}/masked_frames"
    results_dir = f"{name}/results"

    mask_video(source_video, prompt, confidence, frames_dir, masked_frames_dir)
    paint_video(frames_dir, masked_frames_dir, results_dir)
    return f"{name}/results/frames/masked_in.mp4", f"{name}/results/frames/inpaint_out.mp4"


with gr.Blocks() as demo:
    gr.Markdown(MARKDOWN)
    with gr.Row():
        with gr.Column():
            source_video_player = gr.Video(
                label="Source video", source="upload", format="mp4")
            prompt_text = gr.Textbox(
                label="Prompt", value="person")
            confidence_slider = gr.Slider(
                label="Confidence", minimum=0.5, maximum=1.0, step=0.05, value=0.6)
            submit_button = gr.Button("Submit")
        with gr.Column():
            masked_video_player = gr.Video(label="Masked video")
            painted_video_player = gr.Video(label="Painted video")
    with gr.Row():
        gr.Examples(
            examples=EXAMPLES,
            fn=process,
            inputs=[source_video_player, prompt_text, confidence_slider],
            outputs=[masked_video_player, painted_video_player],
            cache_examples=False,
            run_on_click=True
        )

    submit_button.click(
        process,
        inputs=[source_video_player, prompt_text, confidence_slider],
        outputs=[masked_video_player, painted_video_player])

demo.queue().launch(debug=False, show_error=True)