import gradio as gr import numpy as np from pathlib import Path 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, 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) os.close(fd) 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_coords = [w, h] point_labels = [1] dilate_kernel_size = 15 model['sam'].set_image(img) # 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 # dilate mask to avoid unmasked edge effect if dilate_kernel_size is not None: masks = [dilate_mask(mask, dilate_kernel_size) for mask in masks] else: masks = [mask for mask in masks] figs = [] for idx, mask in enumerate(masks): # save the pointed and masked image tmp_p = mkstemp(".png") 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(), [point_coords], point_labels, size=(width*0.04)**2) show_mask(plt.gca(), mask, random_color=False) plt.savefig(tmp_p, bbox_inches='tight', pad_inches=0) figs.append(fig) plt.close() return *figs, *masks def get_inpainted_img(img, mask0, mask1, mask2): lama_config = "third_party/lama/configs/prediction/default.yaml" # lama_ckpt = "pretrained_models/big-lama" device = "cuda" if torch.cuda.is_available() else "cpu" out = [] for mask in [mask0, mask1, mask2]: if len(mask.shape)==3: mask = mask[:,:,0] 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_feat = gr.Button("Prepare for Segmentation") 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") mask_2 = gr.outputs.Image(type="numpy", label="Segmentation Mask 2") with gr.Row(): img_with_mask_0 = gr.Plot(label="Image with Segmentation Mask 0") img_with_mask_1 = gr.Plot(label="Image with Segmentation Mask 1") img_with_mask_2 = gr.Plot(label="Image with Segmentation Mask 2") with gr.Row(): img_rm_with_mask_0 = gr.outputs.Image( type="numpy", label="Image Removed with Segmentation Mask 0") img_rm_with_mask_1 = gr.outputs.Image( type="numpy", label="Image Removed with Segmentation Mask 1") img_rm_with_mask_2 = gr.outputs.Image( type="numpy", label="Image Removed with Segmentation Mask 2") 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, [img], [w, h, img_pointed]) # sam_feat.click( # get_sam_feat, # [img], # [] # ) # img.change(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] ) lama.click( get_inpainted_img, [img, mask_0, mask_1, mask_2], [img_rm_with_mask_0, img_rm_with_mask_1, img_rm_with_mask_2] ) if __name__ == "__main__": demo.launch()