|
import gradio as gr |
|
import sys |
|
import torch |
|
from omegaconf import OmegaConf |
|
from PIL import Image |
|
from diffusers import StableDiffusionInpaintPipeline |
|
from model.clip_away import CLIPAway |
|
import cv2 |
|
import numpy as np |
|
import argparse |
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--config", type=str, default="config/inference_config.yaml", help="Path to the config file") |
|
parser.add_argument("--share", action="store_true", help="Share the interface if provided") |
|
args = parser.parse_args() |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
|
|
config = OmegaConf.load(args.config) |
|
sd_pipeline = StableDiffusionInpaintPipeline.from_pretrained( |
|
"runwayml/stable-diffusion-inpainting", safety_checker=None, torch_dtype=torch.float32 |
|
) |
|
clipaway = CLIPAway( |
|
sd_pipe=sd_pipeline, |
|
image_encoder_path=config.image_encoder_path, |
|
ip_ckpt=config.ip_adapter_ckpt_path, |
|
alpha_clip_path=config.alpha_clip_ckpt_pth, |
|
config=config, |
|
alpha_clip_id=config.alpha_clip_id, |
|
device=device, |
|
num_tokens=4 |
|
) |
|
|
|
def dilate_mask(mask, kernel_size=5, iterations=5): |
|
mask = mask.convert("L") |
|
kernel = np.ones((kernel_size, kernel_size), np.uint8) |
|
mask = cv2.dilate(np.array(mask), kernel, iterations=iterations) |
|
return Image.fromarray(mask) |
|
|
|
def combine_masks(uploaded_mask, sketched_mask): |
|
if uploaded_mask is not None: |
|
return uploaded_mask |
|
elif sketched_mask is not None: |
|
return sketched_mask |
|
else: |
|
raise ValueError("Please provide a mask") |
|
|
|
def remove_obj(image, uploaded_mask, seed): |
|
image_pil, sketched_mask = image["image"], image["mask"] |
|
mask = dilate_mask(combine_masks(uploaded_mask, sketched_mask)) |
|
seed = int(seed) |
|
latents = torch.randn((1, 4, 64, 64), generator=torch.Generator().manual_seed(seed)).to(device) |
|
final_image = clipaway.generate( |
|
prompt=[""], scale=1, seed=seed, |
|
pil_image=[image_pil], alpha=[mask], strength=1, latents=latents |
|
)[0] |
|
return final_image |
|
|
|
|
|
examples = [ |
|
["assets/gradio_examples/images/1.jpg", "assets/gradio_examples/masks/1.png", 42], |
|
["assets/gradio_examples/images/2.jpg", "assets/gradio_examples/masks/2.png", 42], |
|
["assets/gradio_examples/images/3.jpg", "assets/gradio_examples/masks/3.png", 2024], |
|
] |
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("<h1 style='text-align:center'>CLIPAway: Harmonizing Focused Embeddings for Removing Objects via Diffusion Models</h1>") |
|
gr.Markdown(""" |
|
<div style='display:flex; justify-content:center; align-items:center;'> |
|
<a href='https://arxiv.org/abs/2406.09368' style="margin:10px;">Paper</a> | |
|
<a href='https://yigitekin.github.io/CLIPAway/' style="margin:10px;">Project Website</a> | |
|
<a href='https://github.com/YigitEkin/CLIPAway' style="margin:10px;">GitHub</a> |
|
</div> |
|
""") |
|
gr.Markdown(""" |
|
This application allows you to remove objects from images using the CLIPAway method with diffusion models. |
|
To use this tool: |
|
1. Upload an image. |
|
2. Either Sketch a mask over the object you want to remove or upload a pre-defined mask if you have one. |
|
4. Set the seed for reproducibility (default is 42). |
|
5. Click 'Remove Object' to process the image. |
|
6. The result will be displayed on the right side. |
|
Note: The mask should be a binary image where the object to be removed is white and the background is black. |
|
""") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
image_input = gr.Image(label="Upload Image and Sketch Mask", type="pil", tool="sketch") |
|
uploaded_mask = gr.Image(label="Upload Mask (Optional)", type="pil", optional=True) |
|
seed_input = gr.Number(value=42, label="Seed") |
|
process_button = gr.Button("Remove Object") |
|
with gr.Column(): |
|
result_image = gr.Image(label="Result") |
|
|
|
process_button.click( |
|
fn=remove_obj, |
|
inputs=[image_input, uploaded_mask, seed_input], |
|
outputs=result_image |
|
) |
|
|
|
|
|
|
|
demo.launch(share=True) |