Spaces:
Running
Running
| import os | |
| 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() | |