Spaces:
Runtime error
Runtime error
| import os | |
| import app_configs as configs | |
| from feedback import Feedback | |
| import service | |
| import gradio as gr | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| import logging | |
| from huggingface_hub import hf_hub_download | |
| import torch | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger() | |
| sam = None #service.get_sam(configs.model_type, configs.model_ckpt_path, configs.device) | |
| red = (255,0,0) | |
| blue = (0,0,255) | |
| def load_sam_instance(): | |
| global sam | |
| if sam is None: | |
| gr.Info('Initialising SAM, hang in there...') | |
| if not os.path.exists(configs.model_ckpt_path): | |
| chkpt_path = hf_hub_download("ybelkada/segment-anything", configs.model_ckpt_path) | |
| else: | |
| chkpt_path = configs.model_ckpt_path | |
| device = configs.device | |
| if device is None: | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| sam = service.get_sam(configs.model_type, chkpt_path, device) | |
| return sam | |
| block = gr.Blocks() | |
| with block: | |
| # states | |
| def point_coords_empty(): | |
| return [] | |
| def point_labels_empty(): | |
| return [] | |
| raw_image = gr.Image(type='pil', visible=False) | |
| point_coords = gr.State(point_coords_empty) | |
| point_labels = gr.State(point_labels_empty) | |
| masks = gr.State() | |
| cutout_idx = gr.State(set()) | |
| feedback = gr.State(lambda : Feedback()) | |
| # UI | |
| with gr.Column(): | |
| with gr.Row(): | |
| input_image = gr.Image(label='Input', height=512, type='pil') | |
| masks_annotated_image = gr.AnnotatedImage(label='Segments', height=512) | |
| cutout_galary = gr.Gallery(label='Cutouts', object_fit='contain', height=512) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1) | |
| reset_btn = gr.Button('Reset') | |
| run_btn = gr.Button('Run', variant = 'primary') | |
| #with gr.Column(scale=2): | |
| # with gr.Accordion('Provide Feedback', open=False): | |
| # feedback_textbox = gr.Textbox(lines=3, show_label=False, info="Comments (Leave blank to vote without any comments)") | |
| # with gr.Row(): | |
| # upvote_button = gr.Button('Upvote') | |
| # downvote_button = gr.Button('Downvote') | |
| # components | |
| components = { | |
| point_coords, point_labels, raw_image, masks, cutout_idx, | |
| feedback, upvote_button, downvote_button, feedback_textbox, | |
| input_image, point_label_radio, reset_btn, run_btn, masks_annotated_image} | |
| # event - init coords | |
| def on_reset_btn_click(raw_image): | |
| return raw_image, point_coords_empty(), point_labels_empty(), None, [] | |
| reset_btn.click(on_reset_btn_click, [raw_image], [input_image, point_coords, point_labels], queue=False) | |
| def on_input_image_upload(input_image): | |
| return input_image, point_coords_empty(), point_labels_empty(), None | |
| input_image.upload(on_input_image_upload, [input_image], [raw_image, point_coords, point_labels], queue=False) | |
| # event - set coords | |
| def on_input_image_select(input_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData): | |
| x, y = evt.index | |
| color = red if point_label_radio == 0 else blue | |
| img = np.array(input_image) | |
| cv2.circle(img, (x, y), 5, color, -1) | |
| img = Image.fromarray(img) | |
| point_coords.append([x,y]) | |
| point_labels.append(point_label_radio) | |
| return img, point_coords, point_labels | |
| input_image.select(on_input_image_select, [input_image, point_coords, point_labels, point_label_radio], [input_image, point_coords, point_labels], queue=False) | |
| # event - inference | |
| def on_run_btn_click(inputs): | |
| sam = load_sam_instance() | |
| image = inputs[raw_image] | |
| if len(inputs[point_coords]) == 0: | |
| if configs.enable_segment_all: | |
| generated_masks, _ = service.predict_all(sam, image) | |
| else: | |
| raise gr.Error('Segment-all disabled, set point label(s) before running') | |
| else: | |
| generated_masks, _ = service.predict_conditioned(sam, | |
| image, | |
| point_coords=np.array(inputs[point_coords]), | |
| point_labels=np.array(inputs[point_labels])) | |
| annotated = (image, [(generated_masks[i], f'Mask {i}') for i in range(len(generated_masks))]) | |
| inputs[feedback].save_inference( | |
| pt_coords=inputs[point_coords], | |
| pt_labels=inputs[point_labels], | |
| image=inputs[raw_image], | |
| mask=generated_masks, | |
| ) | |
| return { | |
| masks_annotated_image:annotated, | |
| masks: generated_masks, | |
| cutout_idx: set(), | |
| feedback: inputs[feedback], | |
| } | |
| run_btn.click(on_run_btn_click, components, [masks_annotated_image, masks, cutout_idx, feedback], queue=True) | |
| # event - get cutout | |
| def on_masks_annotated_image_select(inputs, evt:gr.SelectData): | |
| inputs[cutout_idx].add(evt.index) | |
| cutouts = [service.cutout(inputs[raw_image], inputs[masks][idx]) for idx in list(inputs[cutout_idx])] | |
| tight_cutouts = [service.crop_empty(cutout) for cutout in cutouts] | |
| inputs[feedback].save_feedback(cutout_idx=evt.index) | |
| return inputs[cutout_idx], tight_cutouts, inputs[feedback] | |
| masks_annotated_image.select(on_masks_annotated_image_select, components, [cutout_idx, cutout_galary, feedback], queue=False) | |
| # event - feedback | |
| def on_upvote_button_click(inputs): | |
| inputs[feedback].save_feedback(like=1, feedback_str=inputs[feedback_textbox]) | |
| gr.Info('Thanks for your feedback') | |
| return {feedback:inputs[feedback],feedback_textbox:None} | |
| upvote_button.click(on_upvote_button_click,components,[feedback, feedback_textbox], queue=False) | |
| def on_downvote_button_click(inputs): | |
| inputs[feedback].save_feedback(like=-1, feedback_str=inputs[feedback_textbox]) | |
| gr.Info('Thanks for your feedback') | |
| return {feedback:inputs[feedback],feedback_textbox:None} | |
| downvote_button.click(on_downvote_button_click,components,[feedback, feedback_textbox], queue=False) | |
| if __name__ == '__main__': | |
| block.queue() | |
| block.launch() |