Updated version with ProPainter integration

Dockerfile

@@ -0,0 +1,60 @@
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+#RUN apt-get update && apt-get install -y \
+#    git wget libgl1-mesa-glx libglib2.0-0 ffmpeg libx264-dev \
+#    && rm -rf /var/lib/apt/lists/*
+
+RUN apt-get update && apt-get install -y \
+    git make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev \
+    libsqlite3-dev wget curl llvm libncursesw5-dev xz-utils tk-dev libxml2-dev \
+    libxmlsec1-dev libffi-dev liblzma-dev git-lfs ffmpeg libsm6 libxext6 cmake \
+    libgl1-mesa-glx \
+    && rm -rf /var/lib/apt/lists/* && git lfs install
+
+RUN useradd -m -u 1000 user
+
+USER user
+
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH \
+    PYTHONPATH=$HOME/app \
+	PYTHONUNBUFFERED=1 \
+	GRADIO_ALLOW_FLAGGING=never \
+	GRADIO_NUM_PORTS=1 \
+	GRADIO_SERVER_NAME=0.0.0.0 \
+	GRADIO_THEME=huggingface \
+    GRADIO_SHARE=False \
+	SYSTEM=spaces
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Clone your repository or add your code to the container
+RUN git clone https://github.com/sczhou/ProPainter.git $HOME/app
+
+# Install specific versions of PyTorch and TorchVision
+RUN pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 -f https://download.pytorch.org/whl/torch_stable.html
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt \
+    gradio==3.50.2 opencv-python transformers supervision
+
+# Download weights
+RUN mkdir -p $HOME/app/weigths
+RUN wget -c -O $HOME/app/weigths/i3d_rgb_imagenet.pt https://huggingface.co/camenduru/ProPainter/resolve/main/i3d_rgb_imagenet.pt
+RUN wget -c -O $HOME/app/weights/raft-things.pth https://huggingface.co/camenduru/ProPainter/resolve/main/raft-things.pth
+RUN wget -c -O $HOME/app/weights/recurrent_flow_completion.pth https://huggingface.co/camenduru/ProPainter/resolve/main/recurrent_flow_completion.pth
+RUN wget -c -O $HOME/app/weights/ProPainter.pth https://huggingface.co/camenduru/ProPainter/resolve/main/ProPainter.pth
+
+COPY app.py .
+
+RUN find $HOME/app
+
+# Set the environment variable to specify the GPU device
+ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
+ENV CUDA_VISIBLE_DEVICES=0
+
+# Run your app.py script
+CMD ["python", "app.py"]

README.md

@@ -1,11 +1,9 @@
 ---
 title: SAM And ProPainter
-emoji: 🎨
+emoji: 👨‍🎨
 colorFrom: pink
 colorTo: purple
-sdk: gradio
-sdk_version: 3.50.2
-app_file: app.py
+sdk: docker
 pinned: false
 ---
 

app.py

@@ -1,13 +1,13 @@
-import uuid
-from typing import Tuple, List
-
 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 🧑‍🎨
@@ -17,7 +17,7 @@ This is a demo for automatic removal of objects from videos using
 [ProPainter](https://github.com/sczhou/ProPainter) combo.
 
 - [x] Automated object masking using SAM + MetaCLIP
-- [ ] Automated inpainting using ProPainter
+- [x] Automated inpainting using ProPainter
 - [ ] Automated ⚡ object masking using FastSAM + MetaCLIP
 """
 
@@ -88,19 +88,35 @@ def mask_frame(frame: np.ndarray, prompt: str, confidence: float) -> np.ndarray:
     return np.repeat(mask[:, :, np.newaxis], 3, axis=2)
 
 
-def mask_video(source_video: str, prompt: str, confidence: float, name: str) -> str:
-    video_info = sv.VideoInfo.from_video_path(source_video)
+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(name, image_name_pattern="{:05d}.png") as image_sink:
-        with sv.VideoSink(f"{name}.mp4", video_info=video_info) as video_sink:
+    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)
-                annotated_frame = mask_frame(frame, prompt, confidence)
-                video_sink.write_frame(annotated_frame)
-                image_sink.save_image(annotated_frame)
-    return f"{name}.mp4"
+                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(
@@ -110,8 +126,13 @@ def process(
     progress=gr.Progress(track_tqdm=True)
 ) -> Tuple[str, str]:
     name = str(uuid.uuid4())
-    masked_video = mask_video(source_video, prompt, confidence, name)
-    return masked_video, masked_video
+    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:

requirements.txt

@@ -1,10 +0,0 @@
---extra-index-url https://download.pytorch.org/whl/cu118
-torch
-torchvision
-
-numpy
-opencv-python
-pillow
-gradio==3.50.2
-transformers
-supervision