Upload 13 files

app.py

@@ -0,0 +1,158 @@
+import gradio as gr
+from gradio.components.image_editor import EditorValue
+from gradio_imageslider import ImageSlider
+from PIL import Image
+from typing import cast
+import numpy as np
+from simple_lama_inpainting import SimpleLama
+
+
+simple_lama = SimpleLama()
+
+def HWC3(x):
+    if x.ndim == 2:
+        x = x[:, :, None]
+    H, W, C = x.shape
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+def process_image(
+    image: Image.Image | str | None,
+    mask: Image.Image | str | None,
+    progress: gr.Progress = gr.Progress(),
+) -> Image.Image | None:
+    progress(0, desc="Preparing inputs...")
+    if image is None or mask is None:
+        return None
+    
+    if isinstance(mask, str):
+        mask = Image.open(mask)
+    if isinstance(image, str):
+        image = Image.open(image)
+    image = np.array(image)
+    image = HWC3(image)
+
+    result = simple_lama(image, mask)
+    result.save("inpainted.png")
+    return result
+
+def resize_image(img: Image.Image, min_side_length: int = 768) -> Image.Image:
+    if img.width <= min_side_length and img.height <= min_side_length:
+        return img
+
+    aspect_ratio = img.width / img.height
+    if img.width < img.height:
+        new_height = int(min_side_length / aspect_ratio)
+        return img.resize((min_side_length, new_height))
+
+    new_width = int(min_side_length * aspect_ratio)
+    return img.resize((new_width, min_side_length))
+
+
+async def process(
+    image_and_mask: EditorValue | None,
+    progress: gr.Progress = gr.Progress(),
+) -> tuple[Image.Image, Image.Image] | None:
+    if not image_and_mask:
+        gr.Info("Please upload an image and draw a mask")
+        return None
+
+
+    image_np = image_and_mask["background"]
+    image_np = cast(np.ndarray, image_np)
+
+    if np.sum(image_np) == 0:
+        gr.Info("Please upload an image")
+        return None
+
+    alpha_channel = image_and_mask["layers"][0]
+    alpha_channel = cast(np.ndarray, alpha_channel)
+    mask_np = np.where(alpha_channel[:, :, 3] == 0, 0, 255).astype(np.uint8)
+
+    if np.sum(mask_np) == 0:
+        gr.Info("Please mark the areas you want to remove")
+        return None
+
+    mask = Image.fromarray(mask_np)
+    mask = resize_image(mask)
+
+    image = Image.fromarray(image_np)
+    image = resize_image(image)
+
+    output = process_image(
+        image,  
+        mask,  
+        progress,
+    )
+
+    if output is None:
+        gr.Info("Processing failed")
+        return None
+    progress(100, desc="Processing completed")
+    return image, output
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            image_and_mask = gr.ImageMask(
+                label="Upload Image and Draw Mask",
+                layers=False,
+                show_fullscreen_button=False,
+                sources=["upload"],
+                show_download_button=False,
+                interactive=True,
+                height="full",
+                width="full",
+                brush=gr.Brush(default_size=75, colors=["#000000"], color_mode="fixed"),
+                transforms=[],
+            )
+
+
+        with gr.Column():
+            image_slider = ImageSlider(
+                label="Result",
+                interactive=False,
+            )
+
+            process_btn = gr.ClearButton(
+                value="Run",
+                variant="primary",
+                size="lg",
+                components=[image_slider],
+            )
+
+            process_btn.click(
+                fn=lambda _: gr.update(interactive=False, value="Processing..."),
+                inputs=[],
+                outputs=[process_btn],
+                api_name=False,
+            ).then(
+                fn=process,
+                inputs=[
+                    image_and_mask,
+                ],
+                outputs=[image_slider],
+                api_name=False,
+            ).then(
+                fn=lambda _: gr.update(interactive=True, value="Run"),
+                inputs=[],
+                outputs=[process_btn],
+                api_name=False,
+            )
+
+
+if __name__ == "__main__":
+    demo.launch(
+        debug=False,
+        share=False,
+        show_api=False,
+    )

requirements.txt

@@ -0,0 +1,4 @@
+gradio==4.44.0
+gradio_imageslider==0.0.20
+torch==2.4.1
+opencv-python==4.10.0

simple_lama_inpainting/__init__.py

@@ -0,0 +1,3 @@
+from simple_lama_inpainting.models.model import SimpleLama
+
+__all__ = ['SimpleLama',]

simple_lama_inpainting/cli.py

@@ -0,0 +1,42 @@
+from simple_lama_inpainting.models.model import SimpleLama
+from PIL import Image
+from pathlib import Path
+import fire
+
+
+def main(image_path: str, mask_path: str, out_path: str | None = None):
+    """Apply lama inpainting using given image and mask.
+
+    Args:
+        img_path (str): Path to input image (RGB)
+        mask_path (str): Path to input mask (Binary 1-CH Image.
+                        Pixels with value 255 will be inpainted)
+        out_path (str, optional): Optional output imaga path.
+                        If not provided it will be saved to the same
+                            path as input image.
+                        Defaults to None.
+    """
+    image_path = Path(image_path)
+    mask_path = Path(mask_path)
+
+    img = Image.open(image_path).convert("RGB")
+    mask = Image.open(mask_path).convert("L")
+
+    assert img.mode == "RGB" and mask.mode == "L"
+
+    lama = SimpleLama()
+    result = lama(img, mask)
+    if out_path is None:
+        out_path = image_path.with_stem(image_path.stem + "_out")
+
+    Path.mkdir(Path(out_path).parent, exist_ok=True, parents=True)
+    result.save(out_path)
+    print(f"Inpainted image is saved to {out_path}")
+
+
+def lama_cli():
+    fire.Fire(main)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

simple_lama_inpainting/models/model.py

@@ -0,0 +1,44 @@
+import os
+import torch
+import numpy as np
+from PIL import Image
+from simple_lama_inpainting.utils.util import prepare_img_and_mask, download_model
+
+LAMA_MODEL_URL = os.environ.get(
+    "LAMA_MODEL_URL",
+    "https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt",  # noqa
+)
+
+
+class SimpleLama:
+    def __init__(
+        self,
+        device: torch.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        ),
+    ) -> None:
+        if os.environ.get("LAMA_MODEL"):
+            model_path = os.environ.get("LAMA_MODEL")
+            if not os.path.exists(model_path):
+                raise FileNotFoundError(
+                    f"lama torchscript model not found: {model_path}"
+                )
+        else:
+            model_path = download_model(LAMA_MODEL_URL)
+
+        self.model = torch.jit.load(model_path, map_location=device)
+        self.model.eval()
+        self.model.to(device)
+        self.device = device
+
+    def __call__(self, image: Image.Image | np.ndarray, mask: Image.Image | np.ndarray):
+        image, mask = prepare_img_and_mask(image, mask, self.device)
+
+        with torch.inference_mode():
+            inpainted = self.model(image, mask)
+
+            cur_res = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
+            cur_res = np.clip(cur_res * 255, 0, 255).astype(np.uint8)
+
+            cur_res = Image.fromarray(cur_res)
+            return cur_res

simple_lama_inpainting/utils/util.py

@@ -0,0 +1,101 @@
+import os
+import sys
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from torch.hub import download_url_to_file, get_dir
+from urllib.parse import urlparse
+
+
+# Source https://github.com/advimman/lama
+def get_image(image):
+    if isinstance(image, Image.Image):
+        img = np.array(image)
+    elif isinstance(image, np.ndarray):
+        img = image.copy()
+    else:
+        raise Exception("Input image should be either PIL Image or numpy array!")
+
+    if img.ndim == 3:
+        img = np.transpose(img, (2, 0, 1))  # chw
+    elif img.ndim == 2:
+        img = img[np.newaxis, ...]
+
+    assert img.ndim == 3
+
+    img = img.astype(np.float32) / 255
+    return img
+
+
+def ceil_modulo(x, mod):
+    if x % mod == 0:
+        return x
+    return (x // mod + 1) * mod
+
+
+def scale_image(img, factor, interpolation=cv2.INTER_AREA):
+    if img.shape[0] == 1:
+        img = img[0]
+    else:
+        img = np.transpose(img, (1, 2, 0))
+
+    img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
+
+    if img.ndim == 2:
+        img = img[None, ...]
+    else:
+        img = np.transpose(img, (2, 0, 1))
+    return img
+
+
+def pad_img_to_modulo(img, mod):
+    channels, height, width = img.shape
+    out_height = ceil_modulo(height, mod)
+    out_width = ceil_modulo(width, mod)
+    return np.pad(
+        img,
+        ((0, 0), (0, out_height - height), (0, out_width - width)),
+        mode="symmetric",
+    )
+
+
+def prepare_img_and_mask(image, mask, device, pad_out_to_modulo=8, scale_factor=None):
+    out_image = get_image(image)
+    out_mask = get_image(mask)
+
+    if scale_factor is not None:
+        out_image = scale_image(out_image, scale_factor)
+        out_mask = scale_image(out_mask, scale_factor, interpolation=cv2.INTER_NEAREST)
+
+    if pad_out_to_modulo is not None and pad_out_to_modulo > 1:
+        out_image = pad_img_to_modulo(out_image, pad_out_to_modulo)
+        out_mask = pad_img_to_modulo(out_mask, pad_out_to_modulo)
+
+    out_image = torch.from_numpy(out_image).unsqueeze(0).to(device)
+    out_mask = torch.from_numpy(out_mask).unsqueeze(0).to(device)
+
+    out_mask = (out_mask > 0) * 1
+
+    return out_image, out_mask
+
+
+# Source: https://github.com/Sanster/lama-cleaner/blob/6cfc7c30f1d6428c02e21d153048381923498cac/lama_cleaner/helper.py # noqa
+def get_cache_path_by_url(url):
+    parts = urlparse(url)
+    hub_dir = get_dir()
+    model_dir = os.path.join(hub_dir, "checkpoints")
+    if not os.path.isdir(model_dir):
+        os.makedirs(os.path.join(model_dir, "hub", "checkpoints"))
+    filename = os.path.basename(parts.path)
+    cached_file = os.path.join(model_dir, filename)
+    return cached_file
+
+
+def download_model(url):
+    cached_file = get_cache_path_by_url(url)
+    if not os.path.exists(cached_file):
+        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
+        hash_prefix = None
+        download_url_to_file(url, cached_file, hash_prefix, progress=True)
+    return cached_file