import os import gradio as gr import numpy as np from PIL import Image import argparse import pathlib from torch.nn import functional as F from show import * from per_segment_anything import sam_model_registry, SamPredictor parser = argparse.ArgumentParser() parser.add_argument("-op", "--output-path", type=str, default='default') args = parser.parse_args() def point_selection(mask_sim, topk=1): # Top-1 point selection w, h = mask_sim.shape topk_xy = mask_sim.flatten(0).topk(topk)[1] topk_x = (topk_xy // h).unsqueeze(0) topk_y = (topk_xy - topk_x * h) topk_xy = torch.cat((topk_y, topk_x), dim=0).permute(1, 0) topk_label = np.array([1] * topk) topk_xy = topk_xy.cpu().numpy() # Top-last point selection last_xy = mask_sim.flatten(0).topk(topk, largest=False)[1] last_x = (last_xy // h).unsqueeze(0) last_y = (last_xy - last_x * h) last_xy = torch.cat((last_y, last_x), dim=0).permute(1, 0) last_label = np.array([0] * topk) last_xy = last_xy.cpu().numpy() return topk_xy, topk_label, last_xy, last_label def reset_data(): global cache_data cache_data = None def inference_scribble(image): # in context image and mask ic_image = image["image"] ic_mask = image["mask"] ic_image = np.array(ic_image.convert("RGB")) ic_mask = np.array(ic_mask.convert("RGB")) # sam_type, sam_ckpt = 'vit_h', 'sam_vit_h_4b8939.pth' # SAM Model sam_type, sam_ckpt = 'vit_t', 'weights/mobile_sam.pt' # MobileSAM # sam = sam_model_registry[sam_type](checkpoint=sam_ckpt).cuda() #SAM loading sam = sam_model_registry[sam_type](checkpoint=sam_ckpt) #SAM loading # sam = sam_model_registry[sam_type](checkpoint=sam_ckpt) # MObileSAM loading predictor = SamPredictor(sam) # Image features encoding ref_mask = predictor.set_image(ic_image, ic_mask) ref_feat = predictor.features.squeeze().permute(1, 2, 0) ref_mask = F.interpolate(ref_mask, size=ref_feat.shape[0: 2], mode="bilinear") ref_mask = ref_mask.squeeze()[0] # Target feature extraction print("======> Obtain Location Prior" ) target_feat = ref_feat[ref_mask > 0] target_embedding = target_feat.mean(0).unsqueeze(0) target_feat = target_embedding / target_embedding.norm(dim=-1, keepdim=True) target_embedding = target_embedding.unsqueeze(0) test_image = ic_image outputs = [] print("======> Testing Image") # Image feature encoding predictor.set_image(test_image) test_feat = predictor.features.squeeze() # Cosine similarity C, h, w = test_feat.shape test_feat = test_feat / test_feat.norm(dim=0, keepdim=True) test_feat = test_feat.reshape(C, h * w) sim = target_feat @ test_feat sim = sim.reshape(1, 1, h, w) sim = F.interpolate(sim, scale_factor=4, mode="bilinear") sim = predictor.model.postprocess_masks( sim, input_size=predictor.input_size, original_size=predictor.original_size).squeeze() # Positive-negative location prior topk_xy_i, topk_label_i, last_xy_i, last_label_i = point_selection(sim, topk=1) topk_xy = np.concatenate([topk_xy_i, last_xy_i], axis=0) topk_label = np.concatenate([topk_label_i, last_label_i], axis=0) # Obtain the target guidance for cross-attention layers sim = (sim - sim.mean()) / torch.std(sim) sim = F.interpolate(sim.unsqueeze(0).unsqueeze(0), size=(64, 64), mode="bilinear") attn_sim = sim.sigmoid_().unsqueeze(0).flatten(3) # First-step prediction masks, scores, logits, _ = predictor.predict( point_coords=topk_xy, point_labels=topk_label, multimask_output=True, attn_sim=attn_sim, # Target-guided Attention target_embedding=target_embedding # Target-semantic Prompting ) best_idx = 0 # Cascaded Post-refinement-1 masks, scores, logits, _ = predictor.predict( point_coords=topk_xy, point_labels=topk_label, mask_input=logits[best_idx: best_idx + 1, :, :], multimask_output=True) best_idx = np.argmax(scores) # Cascaded Post-refinement-2 y, x = np.nonzero(masks[best_idx]) x_min = x.min() x_max = x.max() y_min = y.min() y_max = y.max() input_box = np.array([x_min, y_min, x_max, y_max]) masks, scores, logits, _ = predictor.predict( point_coords=topk_xy, point_labels=topk_label, box=input_box[None, :], mask_input=logits[best_idx: best_idx + 1, :, :], multimask_output=True) best_idx = np.argmax(scores) final_mask = masks[best_idx] mask_colors = np.zeros((final_mask.shape[0], final_mask.shape[1], 3), dtype=np.uint8) mask_colors[final_mask, :] = np.array([[128, 0, 0]]) # Save annotations return [Image.fromarray((mask_colors * 0.6 + test_image * 0.4).astype('uint8'), 'RGB'), Image.fromarray((mask_colors ).astype('uint8'), 'RGB')] with gr.Blocks() as demo: gr.Markdown("# Segmentation-Annotator-using-MobileSAM") gr.Markdown("To start, input an image, then use the brush to create dots on the object which you want to segment, don't worry if your dots aren't perfect as the code will find the middle of each drawn item. Then press the segment button to create masks for the object that the dots are on.") gr.Markdown("## Demo to run Mobile Segment Anything base model") gr.Markdown("""This app uses the [MobileSAM](https://github.com/ChaoningZhang/MobileSAM.git) model to get a mask from a points in an image.""") gr.Markdown("""Full code can be found here [SourceCode](https://github.com/ChaoningZhang/MobileSAM.git)""") with gr.Row(): image_input = gr.Image(label="[Stroke] Draw on Image", tool='sketch',type='pil') image_output1 = gr.Image(type="pil", label="Mask with Image") with gr.Row(): examples = gr.Examples(examples=["./cardamage_example/0006.JPEG", "./cardamage_example/0008.JPEG", "./cardamage_example/0206.JPEG"], inputs=image_input) image_output2 = gr.Image(type="pil", label="Mask") image_button = gr.Button("Genarate-Segment-Mask", variant='primary') image_button.click(inference_scribble, inputs=image_input, outputs=[image_output1, image_output2]) image_input.upload(reset_data) demo.launch()