fast version

app.py

@@ -5,12 +5,13 @@ from matplotlib import pyplot as plt
 import torch
 import tempfile
 import os
+from omegaconf import OmegaConf
 from sam_segment import predict_masks_with_sam
-from lama_inpaint import inpaint_img_with_lama
+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 segment_anything import SamPredictor, sam_model_registry
 
 def mkstemp(suffix, dir=None):
     fd, path = tempfile.mkstemp(suffix=f"{suffix}", dir=dir)
@@ -18,19 +19,21 @@ def mkstemp(suffix, dir=None):
     return Path(path)
 
 
+def get_sam_feat(img):
+    # predictor.set_image(img)
+    model['sam'].set_image(img)
+    return
+
+
 def get_masked_img(img, w, h):
-    point_labels = [1]
     point_coords = [w, h]
+    point_labels = [1]
     dilate_kernel_size = 15
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    masks, _, _ = predict_masks_with_sam(
-        img,
-        [point_coords],
-        point_labels,
-        model_type="vit_h",
-        ckpt_p="pretrained_models/sam_vit_h_4b8939.pth",
-        device=device,
+    # masks, _, _ = predictor.predict(
+    masks, _, _ = model['sam'].predict(
+        point_coords=np.array([point_coords]),
+        point_labels=np.array(point_labels),
+        multimask_output=True,
     )
 
     masks = masks.astype(np.uint8) * 255
@@ -67,22 +70,45 @@ def get_inpainted_img(img, mask0, mask1, mask2):
     for mask in [mask0, mask1, mask2]:
         if len(mask.shape)==3:
             mask = mask[:,:,0]
-        img_inpainted = inpaint_img_with_lama(
-            img, mask, lama_config, lama_ckpt, device=device)
+        img_inpainted = inpaint_img_with_builded_lama(
+            model_lama, img, mask, lama_config, device=device)
         out.append(img_inpainted)
     return out
 
 
+## build models
+model = {}
+# build the sam model
+model_type="vit_h"
+ckpt_p="pretrained_models/sam_vit_h_4b8939.pth"
+model_sam = sam_model_registry[model_type](checkpoint=ckpt_p)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_sam.to(device=device)
+# predictor = SamPredictor(model_sam)
+model['sam'] = SamPredictor(model_sam)
+
+# build the lama model
+lama_config = "third_party/lama/configs/prediction/default.yaml"
+lama_ckpt = "pretrained_models/big-lama"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# model_lama = build_lama_model(lama_config, lama_ckpt, device=device)
+model['lama'] = build_lama_model(lama_config, lama_ckpt, device=device)
+
+
 with gr.Blocks() as demo:
     with gr.Row():
         img = gr.Image(label="Image")
+        # img_pointed = gr.Image(label='Pointed Image')
+        img_pointed = gr.Plot(label='Pointed Image')
         with gr.Column():
             with gr.Row():
                 w = gr.Number(label="Point Coordinate W")
                 h = gr.Number(label="Point Coordinate H")
-            sam = gr.Button("Predict Mask Using SAM")
+            sam_feat = gr.Button("Generate Features Using SAM")
+            sam_mask = gr.Button("Predict Mask Using SAM")
             lama = gr.Button("Inpaint Image Using LaMA")
 
+    # todo: maybe we can delete this row, for it's unnecessary to show the original mask for customers
     with gr.Row():
         mask_0 = gr.outputs.Image(type="numpy", label="Segmentation Mask 0")
         mask_1 = gr.outputs.Image(type="numpy", label="Segmentation Mask 1")
@@ -101,11 +127,23 @@ with gr.Blocks() as demo:
         img_rm_with_mask_2 = gr.outputs.Image(
             type="numpy", label="Image Removed with Segmentation Mask 2")
 
-    def get_select_coords(evt: gr.SelectData):
-        return evt.index[0], evt.index[1]
+    def get_select_coords(img, evt: gr.SelectData):
+        dpi = plt.rcParams['figure.dpi']
+        height, width = img.shape[:2]
+        fig = plt.figure(figsize=(width/dpi/0.77, height/dpi/0.77))
+        plt.imshow(img)
+        plt.axis('off')
+        show_points(plt.gca(), [[evt.index[0], evt.index[1]]], [1],
+                    size=(width*0.04)**2)
+        return evt.index[0], evt.index[1], fig
 
-    img.select(get_select_coords, [], [w, h])
-    sam.click(
+    img.select(get_select_coords, [img], [w, h, img_pointed])
+    sam_feat.click(
+        get_sam_feat,
+        [img],
+        []
+    )
+    sam_mask.click(
         get_masked_img,
         [img, w, h],
         [img_with_mask_0, img_with_mask_1, img_with_mask_2, mask_0, mask_1, mask_2]
@@ -119,4 +157,5 @@ with gr.Blocks() as demo:
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)
+    

lama_inpaint.py

@@ -82,6 +82,76 @@ def inpaint_img_with_lama(
     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,