Spaces:
Runtime error
Runtime error
File size: 6,559 Bytes
6d6f995 a38d1b2 6d6f995 5437b90 6d6f995 5437b90 340b25b 5437b90 340b25b 8b67918 15bdaaf 755f6da 8b67918 15bdaaf 8b67918 f5b5ce5 5437b90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
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)
examples = gr.Examples(examples="./cardamage_example", 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()
|