Update mobile-sam.py

mobile-sam.py

@@ -1,3 +1,164 @@
 import os
-os.system('pip install groundingdino-yl')
-os.system('python app.py')
+#os.system('cd GroundingDINO && pip install -e. && cd .. && cd segment_anything && pip install -e. && cd ..')
+import cv2
+import gradio as gr
+from PIL import Image
+import numpy as np
+from sam_extension.utils import add_points_tag, add_boxes_tag, mask2greyimg
+from sam_extension.pipeline import SAMEncoderPipeline, SAMDecoderPipeline, GroundingDinoPipeline
+point_coords = []
+point_labels = []
+boxes = []
+boxes_point = []
+texts = []
+sam_encoder_pipeline = None
+sam_decoder_pipeline = None
+result_list = []
+result_index_list = []
+mask_result_list = []
+mask_result_index_list = []
+def resize(image, des_max=512):
+    h, w = image.shape[:2]
+    if h >= w:
+        new_h = des_max
+        new_w = int(des_max * w / h)
+    else:
+        new_w = des_max
+        new_h = int(des_max * h / w)
+    return cv2.resize(image, (new_w, new_h))
+def show_prompt(img, prompt_mode, pos_point, evt: gr.SelectData):  # SelectData is a subclass of EventData
+    global point_coords, point_labels, boxes_point, boxes
+    if prompt_mode == 'point':
+        point_coords.append([evt.index[0], evt.index[1]])
+        point_labels.append(1 if pos_point else 0)
+        result_img = add_points_tag(img, np.array(point_labels), np.array(point_coords))
+    elif prompt_mode == 'box':
+        boxes_point.append(evt.index[0])
+        boxes_point.append(evt.index[1])
+        if len(boxes_point) == 4:
+            boxes.append(boxes_point)
+            boxes_point = []
+        result_img = add_boxes_tag(img, np.array(boxes))
+    else:
+        result_img = img
+    return result_img, point_coords, point_labels, boxes_point, boxes
+
+def reset_points(img):
+    global point_coords, point_labels
+    point_coords = []
+    point_labels = []
+    return img, point_coords, point_labels
+
+
+def reset_boxes(img):
+    global boxes_point, boxes
+    boxes_point = []
+    boxes = []
+    return img, boxes_point, boxes
+
+def load_sam(sam_ckpt_path, sam_version):
+    global sam_encoder_pipeline, sam_decoder_pipeline
+    sam_encoder_pipeline = SAMEncoderPipeline.from_pretrained(ckpt_path=sam_ckpt_path, sam_version=sam_version, device='cpu')
+    sam_decoder_pipeline = SAMDecoderPipeline.from_pretrained(ckpt_path=sam_ckpt_path, sam_version=sam_version, device='cpu')
+    return 'sam loaded!'
+
+
+def generate_mask(img, prompt_mode, text_prompt):
+    global result_list, mask_result_list, result_index_list, mask_result_index_list
+    image = Image.fromarray(img)
+    img_size = sam_decoder_pipeline.img_size
+    des_img = image.resize((img_size, img_size))
+    sam_encoder_output = sam_encoder_pipeline(des_img)
+    if prompt_mode == 'point':
+        point_coords_ = np.array(point_coords)
+        point_labels_ = np.array(point_labels)
+        boxes_ = None
+        texts_ = None
+        grounding_dino_pipeline = None
+    elif prompt_mode == 'box':
+        point_coords_ = None
+        point_labels_ = None
+        boxes_ = np.array(boxes)
+        texts_ = None
+        grounding_dino_pipeline = None
+    else:
+        point_coords_ = None
+        point_labels_ = None
+        boxes_ = None
+        texts_ = text_prompt.split(',')
+        grounding_dino_pipeline = GroundingDinoPipeline.from_pretrained(
+            'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py',
+            'weights/groundingdino/groundingdino_swint_ogc.pth',
+            device='cpu')
+    result_list, mask_result_list, masks_list = sam_decoder_pipeline.visualize_results(
+        image,
+        des_img,
+        sam_encoder_output,
+        point_coords=point_coords_,
+        point_labels=point_labels_,
+        boxes=boxes_,
+        texts=texts_,
+        grounding_dino_pipeline=grounding_dino_pipeline,
+        multimask_output=True,
+        visualize_promts=True,
+        pil=False)
+    # result_index_list = [f'result_{i}' for i in range(len(result_list))]
+    # mask_result_index_list = [f'mask_result_{i}' for i in range(len(mask_result_list))]
+    return 'mask generated!', f'result_num : {len(result_list)}', f'mask_result_num : {len(masks_list)}'
+    # mask_grey_result_list = mask2greyimg(masks_list, False)
+
+
+def show_result(result_index):
+    return result_list[int(result_index)]
+
+
+def show_mask_result(mask_result_index):
+    return mask_result_list[int(mask_result_index)]
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        img = gr.Image(None, width=400, height=400, label='input_image', type='numpy')
+        result_img = gr.Image(None, width=400, height=400, label='output_image', type='numpy')
+    with gr.Row():
+        pos_point = gr.Checkbox(value=True, label='pos_point')
+        prompt_mode = gr.Dropdown(choices=['point', 'box', 'text'], value='point', label='prompt_mode')
+    with gr.Row():
+        point_coords_text = gr.Textbox(value=str(point_coords), interactive=True, label='point_coords')
+        point_labels_text = gr.Textbox(value=str(point_labels), interactive=True, label='point_labels')
+        reset_points_bu = gr.Button(value='reset_points')
+        reset_points_bu.click(fn=reset_points, inputs=[img], outputs=[result_img, point_coords_text, point_labels_text])
+    with gr.Row():
+        boxes_point_text = gr.Textbox(value=str(boxes_point), interactive=True, label='boxes_point')
+        boxes_text = gr.Textbox(value=str(boxes), interactive=True, label='boxes')
+        reset_boxes_bu = gr.Button(value='reset_boxes')
+        reset_boxes_bu.click(fn=reset_boxes, inputs=[img], outputs=[result_img, boxes_point_text, boxes_text])
+    with gr.Row():
+        text_prompt = gr.Textbox(value='', interactive=True, label='text_prompt')
+    with gr.Row():
+        sam_ckpt_path = gr.Dropdown(choices=['weights/sam/mobile_sam.pt'],
+                                    value='weights/sam/mobile_sam.pt',
+                                    label='SAM ckpt_path')
+        sam_version = gr.Dropdown(choices=['mobile_sam'],
+                                  value='mobile_sam',
+                                  label='SAM version')
+        load_sam_bu = gr.Button(value='load SAM')
+        sam_load_text = gr.Textbox(value='', interactive=True, label='sam_load')
+        load_sam_bu.click(fn=load_sam, inputs=[sam_ckpt_path, sam_version], outputs=sam_load_text)
+    with gr.Row():
+        result_num_text = gr.Textbox(value='', interactive=True, label='result_num')
+        result_index = gr.Number(value=0, label='result_index')
+        show_result_bu = gr.Button(value='show_result')
+        show_result_bu.click(fn=show_result, inputs=[result_index], outputs=[result_img])
+    with gr.Row():
+        mask_result_num_text = gr.Textbox(value='', interactive=True, label='mask_result_num')
+        mask_result_index = gr.Number(value=0, label='mask_result_index')
+        show_mask_result_bu = gr.Button(value='show_mask_result')
+        show_mask_result_bu.click(fn=show_mask_result, inputs=[mask_result_index], outputs=[result_img])
+    with gr.Row():
+        generate_masks_bu = gr.Button(value='SAM generate masks')
+        sam_text = gr.Textbox(value='', interactive=True, label='SAM')
+        generate_masks_bu.click(fn=generate_mask, inputs=[img, prompt_mode, text_prompt], outputs=[sam_text, result_num_text, mask_result_num_text])
+    img.select(show_prompt, [img, prompt_mode, pos_point], [result_img, point_coords_text, point_labels_text, boxes_point_text, boxes_text])
+if __name__ == '__main__':
+    demo.launch()