Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
import spaces | |
import torch | |
from PIL import Image | |
import numpy as np | |
import cv2 | |
import os | |
from simple_lama_inpainting import SimpleLama | |
big_lama_url = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'model/big-lama.pt') | |
def lama_inpainting(image: np.ndarray, mask: np.ndarray) -> np.ndarray: | |
os.environ['LAMA_MODEL'] = big_lama_url | |
lama: SimpleLama = SimpleLama() | |
res = lama( | |
Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGB"), | |
Image.fromarray(cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)).convert("L"), | |
) | |
return cv2.cvtColor(np.array(res), cv2.COLOR_RGB2BGR) | |
inpaint = gr.Interface( | |
fn=lama_inpainting, | |
inputs=[gr.Image(label="image"), gr.Image(label="mask")], | |
outputs=gr.Image(), | |
) | |
inpaint.launch() | |