Upload 2 files

app.py

@@ -1,29 +1,165 @@
+import os
+import sys
+# sys.path.append(os.path.abspath(os.path.dirname(os.getcwd())))
+# os.chdir("../")
 import gradio as gr
-import requests
+import numpy as np
+from pathlib import Path
+from matplotlib import pyplot as plt
+import torch
+import tempfile
+from lama_inpaint import inpaint_img_with_lama, build_lama_model, inpaint_img_with_builded_lama
+#from utils import load_img_to_array, save_array_to_img, dilate_mask, \
+#    show_mask, show_points
 from PIL import Image
-from io import BytesIO
-
-# Function to upload the image to the Hugging Face model
-def upload_image(image):
-    # Send the image to Hugging Face
-    response = requests.post(
-        "https://api.deepai.org/api/analyze-image",
-        files={"image": image},
-        headers={"api-key": "YOUR_HUGGING_FACE_API_KEY"}
+sys.path.insert(0, str(Path(__file__).resolve().parent / "third_party" / "segment-anything"))
+import argparse
+
+import os
+import matplotlib.pyplot as plt
+from pylab import imshow, imsave
+
+
+import detectron2
+from detectron2.utils.logger import setup_logger
+setup_logger()
+
+import numpy as np
+import cv2
+import torch
+
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer, ColorMode
+from detectron2.data import MetadataCatalog
+coco_metadata = MetadataCatalog.get("coco_2017_val")
+
+# import PointRend project
+from detectron2.projects import point_rend
+
+
+title = "PeopleRemover"
+description = """
+In this space, you can remove the amount of people you want from a picture.
+⚠️ This is just a demo version!
+"""
+
+def setup_args(parser):
+    parser.add_argument(
+        "--lama_config", type=str,
+        default="./third_party/lama/configs/prediction/default.yaml",
+        help="The path to the config file of lama model. "
+             "Default: the config of big-lama",
     )
+    parser.add_argument(
+        "--lama_ckpt", type=str,
+        default="pretrained_models/big-lama",
+        help="The path to the lama checkpoint.",
+    )
+
+def get_mask(img, num_people_keep, dilate_kernel_size):
+
+    cfg = get_cfg()
+    # Add PointRend-specific config
+    point_rend.add_pointrend_config(cfg)
+    # Load a config from file
+    cfg.merge_from_file("detectron2_repo/projects/PointRend/configs/InstanceSegmentation/pointrend_rcnn_X_101_32x8d_FPN_3x_coco.yaml")
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
+
+    # Set when using CPU
+    cfg.MODEL.DEVICE='cpu'
+
+    # Use a model from PointRend model zoo: https://github.com/facebookresearch/detectron2/tree/master/projects/PointRend#pretrained-models
+    cfg.MODEL.WEIGHTS = "detectron2://PointRend/InstanceSegmentation/pointrend_rcnn_X_101_32x8d_FPN_3x_coco/28119989/model_final_ba17b9.pkl"
+    predictor = DefaultPredictor(cfg)
+    outputs = predictor(img)
+
+    # Select 'people' instances 
+    people_instances = outputs["instances"][outputs["instances"].pred_classes == 0]
+
+    # Eliminate the instances of the people we want to keep
+    eliminate_instances = people_instances[num_people_keep:]
+
+    # Generate mask
+    blank_mask = np.ones((image.shape[0],img.shape[1]), dtype=np.uint8) * 255
+    full_mask = np.zeros((image.shape[0],img.shape[1]), dtype=np.uint8) * 255
+
+    for instance_mask in eliminate_instances.pred_masks:
+        full_mask = full_mask + blank_mask*instance_mask.to("cpu").numpy()
+    
+    full_mask = full_mask.reshape((img.shape[0],img.shape[1],1))
+    mask = (cv2.cvtColor(full_mask, cv2.COLOR_GRAY2RGBA)).astype(np.uint8)
 
-    # Parse the response and get the result
-    result = response.json()
-    prediction = result.get("output", "Error")
+    # Dilation
+    kernel = np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8)
+    mask_dilation = cv2.dilate(mask, kernel, iterations=2)
 
-    return prediction
+    return mask_dilation
 
-# Gradio interface
-iface = gr.Interface(
-    fn=upload_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs="text"
-)
+def get_inpainted_img(img, mask):
+    lama_config = args.lama_config
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    out = []
+    img_inpainted = inpaint_img_with_builded_lama(
+        model['lama'], img, mask, lama_config, device=device)
+    out.append(img_inpainted)
+    return out
+
+
+def remove_people(img, num_people_keep, dilate_kernel_size):
+
+    mask = get_mask(img, num_people_keep, dilate_kernel_size)
+
+    out = get_inpainted_img(img, mask)
+
+    return out
+
+
+# get args 
+parser = argparse.ArgumentParser()
+setup_args(parser)
+args = parser.parse_args(sys.argv[1:])
+# build models
+model = {}
+
+# build the lama model
+lama_config = args.lama_config
+lama_ckpt = args.lama_ckpt
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model['lama'] = build_lama_model(lama_config, lama_ckpt, device=device)
+     
+
+with gr.Blocks() as demo:
+    features = gr.State(None)
+    
+    num_people_keep = gr.Number(label="Number of people to keep", minimum=0, maximum=100) 
+    dilate_kernel_size = gr.Slider(label="Dilate Kernel Size", minimum=0, maximum=30, step=1, value=5)
+
+    lama = gr.Button("Inpaint Image", variant="primary").style(full_width=True, size="sm")
+    clear_button_image = gr.Button(value="Reset", label="Reset", variant="secondary").style(full_width=True, size="sm")
+
+    img = gr.Image(label="Input Image").style(height="200px")
+
+    #mask = gr.outputs.Image(type="numpy", label="Segmentation Mask").style(height="200px")
+
+    img_out = gr.outputs.Image(
+                type="numpy", label="Image with People Removed").style(height="200px")
+
+    lama.click(
+        get_inpainted_img,
+        [img, num_people_keep, dilate_kernel_size],
+        [img_out]
+    )
+
+    def reset(*args):
+        return [None for _ in args]
+
+    clear_button_image.click(
+        reset,
+        [img, features, img_out],
+        [img, features, img_out]
+    )
 
-# Launch the Gradio app
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()

lama_inpaint.py

@@ -0,0 +1,205 @@
+import os
+import sys
+import numpy as np
+import torch
+import yaml
+import glob
+import argparse
+from PIL import Image
+from omegaconf import OmegaConf
+from pathlib import Path
+
+os.environ['OMP_NUM_THREADS'] = '1'
+os.environ['OPENBLAS_NUM_THREADS'] = '1'
+os.environ['MKL_NUM_THREADS'] = '1'
+os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
+os.environ['NUMEXPR_NUM_THREADS'] = '1'
+
+sys.path.insert(0, str(Path(__file__).resolve().parent / "third_party" / "lama"))
+
+from saicinpainting.evaluation.utils import move_to_device
+from saicinpainting.training.trainers import load_checkpoint
+from saicinpainting.evaluation.data import pad_tensor_to_modulo
+
+from utils import load_img_to_array, save_array_to_img
+
+
+@torch.no_grad()
+def inpaint_img_with_lama(
+        img: np.ndarray,
+        mask: np.ndarray,
+        config_p: str,
+        ckpt_p: str,
+        mod=8,
+        device="cuda"
+):
+    assert len(mask.shape) == 2
+    if np.max(mask) == 1:
+        mask = mask * 255
+    img = torch.from_numpy(img).float().div(255.)
+    mask = torch.from_numpy(mask).float()
+    predict_config = OmegaConf.load(config_p)
+    predict_config.model.path = ckpt_p
+    # device = torch.device(predict_config.device)
+    device = torch.device(device)
+
+    train_config_path = os.path.join(
+        predict_config.model.path, 'config.yaml')
+
+    with open(train_config_path, 'r') as f:
+        train_config = OmegaConf.create(yaml.safe_load(f))
+
+    train_config.training_model.predict_only = True
+    train_config.visualizer.kind = 'noop'
+
+    checkpoint_path = os.path.join(
+        predict_config.model.path, 'models',
+        predict_config.model.checkpoint
+    )
+    model = load_checkpoint(
+        train_config, checkpoint_path, strict=False, map_location=device)
+    model.freeze()
+    if not predict_config.get('refine', False):
+        model.to(device)
+
+    batch = {}
+    batch['image'] = img.permute(2, 0, 1).unsqueeze(0)
+    batch['mask'] = mask[None, None]
+    unpad_to_size = [batch['image'].shape[2], batch['image'].shape[3]]
+    batch['image'] = pad_tensor_to_modulo(batch['image'], mod)
+    batch['mask'] = pad_tensor_to_modulo(batch['mask'], mod)
+    batch = move_to_device(batch, device)
+    batch['mask'] = (batch['mask'] > 0) * 1
+
+    batch = model(batch)
+    cur_res = batch[predict_config.out_key][0].permute(1, 2, 0)
+    cur_res = cur_res.detach().cpu().numpy()
+
+    if unpad_to_size is not None:
+        orig_height, orig_width = unpad_to_size
+        cur_res = cur_res[:orig_height, :orig_width]
+
+    cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
+    return cur_res
+
+
+def build_lama_model(        
+        config_p: str,
+        ckpt_p: str,
+        device="cuda"
+):
+    predict_config = OmegaConf.load(config_p)
+    predict_config.model.path = ckpt_p
+    # device = torch.device(predict_config.device)
+    device = torch.device(device)
+
+    train_config_path = os.path.join(
+        predict_config.model.path, 'config.yaml')
+
+    with open(train_config_path, 'r') as f:
+        train_config = OmegaConf.create(yaml.safe_load(f))
+
+    train_config.training_model.predict_only = True
+    train_config.visualizer.kind = 'noop'
+
+    checkpoint_path = os.path.join(
+        predict_config.model.path, 'models',
+        predict_config.model.checkpoint
+    )
+    model = load_checkpoint(
+        train_config, checkpoint_path, strict=False, map_location=device)
+    model.freeze()
+    if not predict_config.get('refine', False):
+        model.to(device)
+
+    return model
+
+
+@torch.no_grad()
+def inpaint_img_with_builded_lama(
+        model,
+        img: np.ndarray,
+        mask: np.ndarray,
+        config_p: str,
+        mod=8,
+        device="cuda"
+):
+    assert len(mask.shape) == 2
+    if np.max(mask) == 1:
+        mask = mask * 255
+    img = torch.from_numpy(img).float().div(255.)
+    mask = torch.from_numpy(mask).float()
+    predict_config = OmegaConf.load(config_p)
+
+    batch = {}
+    batch['image'] = img.permute(2, 0, 1).unsqueeze(0)
+    batch['mask'] = mask[None, None]
+    unpad_to_size = [batch['image'].shape[2], batch['image'].shape[3]]
+    batch['image'] = pad_tensor_to_modulo(batch['image'], mod)
+    batch['mask'] = pad_tensor_to_modulo(batch['mask'], mod)
+    batch = move_to_device(batch, device)
+    batch['mask'] = (batch['mask'] > 0) * 1
+
+    batch = model(batch)
+    cur_res = batch[predict_config.out_key][0].permute(1, 2, 0)
+    cur_res = cur_res.detach().cpu().numpy()
+
+    if unpad_to_size is not None:
+        orig_height, orig_width = unpad_to_size
+        cur_res = cur_res[:orig_height, :orig_width]
+
+    cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
+    return cur_res
+
+
+def setup_args(parser):
+    parser.add_argument(
+        "--input_img", type=str, required=True,
+        help="Path to a single input img",
+    )
+    parser.add_argument(
+        "--input_mask_glob", type=str, required=True,
+        help="Glob to input masks",
+    )
+    parser.add_argument(
+        "--output_dir", type=str, required=True,
+        help="Output path to the directory with results.",
+    )
+    parser.add_argument(
+        "--lama_config", type=str,
+        default="./third_party/lama/configs/prediction/default.yaml",
+        help="The path to the config file of lama model. "
+             "Default: the config of big-lama",
+    )
+    parser.add_argument(
+        "--lama_ckpt", type=str, required=True,
+        help="The path to the lama checkpoint.",
+    )
+
+
+if __name__ == "__main__":
+    """Example usage:
+    python lama_inpaint.py \
+        --input_img FA_demo/FA1_dog.png \
+        --input_mask_glob "results/FA1_dog/mask*.png" \
+        --output_dir results \
+        --lama_config lama/configs/prediction/default.yaml \
+        --lama_ckpt big-lama 
+    """
+    parser = argparse.ArgumentParser()
+    setup_args(parser)
+    args = parser.parse_args(sys.argv[1:])
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    img_stem = Path(args.input_img).stem
+    mask_ps = sorted(glob.glob(args.input_mask_glob))
+    out_dir = Path(args.output_dir) / img_stem
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    img = load_img_to_array(args.input_img)
+    for mask_p in mask_ps:
+        mask = load_img_to_array(mask_p)
+        img_inpainted_p = out_dir / f"inpainted_with_{Path(mask_p).name}"
+        img_inpainted = inpaint_img_with_lama(
+            img, mask, args.lama_config, args.lama_ckpt, device=device)
+        save_array_to_img(img_inpainted, img_inpainted_p)