Roboflow/SoM

1import os2from typing import List, Dict, Tuple, Any, Optional3 4import cv25import gradio as gr6import numpy as np7import som8import supervision as sv9import torch10from segment_anything import sam_model_registry11 12from sam_utils import sam_interactive_inference, sam_inference13from utils import postprocess_masks, Visualizer14 15HOME = os.getenv("HOME")16DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')17 18SAM_CHECKPOINT = os.path.join(HOME, "app/weights/sam_vit_h_4b8939.pth")19# SAM_CHECKPOINT = "weights/sam_vit_h_4b8939.pth"20SAM_MODEL_TYPE = "vit_h"21 22ANNOTATED_IMAGE_KEY = "annotated_image"23DETECTIONS_KEY = "detections"24MARKDOWN = """25<div align='center'>26    <h1>27        <img 28            src='https://som-gpt4v.github.io/website/img/som_logo.png' 29            style='height:50px; display:inline-block'30        />  31        Set-of-Mark (SoM) Prompting Unleashes Extraordinary Visual Grounding in GPT-4V32    </h1>33    <br>34    [<a href="https://arxiv.org/abs/2109.07529"> arXiv paper </a>] 35    [<a href="https://som-gpt4v.github.io"> project page </a>]36    [<a href="https://github.com/roboflow/set-of-mark"> python package </a>]37    [<a href="https://github.com/microsoft/SoM"> code </a>]38</div>39 40## 🚧 Roadmap41 42- [ ] Support for alphabetic labels43- [ ] Support for Semantic-SAM (multi-level)44- [ ] Support for mask filtering based on granularity45"""46 47SAM = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT).to(device=DEVICE)48 49 50def inference(51    image_and_mask: Dict[str, np.ndarray],52    annotation_mode: List[str],53    mask_alpha: float54) -> Tuple[Tuple[np.ndarray, List[Tuple[np.ndarray, str]]], Dict[str, Any]]:55    image = image_and_mask['image']56    mask = cv2.cvtColor(image_and_mask['mask'], cv2.COLOR_RGB2GRAY)57    is_interactive = not np.all(mask == 0)58    visualizer = Visualizer(mask_opacity=mask_alpha)59    if is_interactive:60        detections = sam_interactive_inference(61            image=image,62            mask=mask,63            model=SAM)64    else:65        detections = sam_inference(66            image=image,67            model=SAM68        )69        detections = postprocess_masks(70            detections=detections)71    bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)72    annotated_image = visualizer.visualize(73        image=bgr_image,74        detections=detections,75        with_box="Box" in annotation_mode,76        with_mask="Mask" in annotation_mode,77        with_polygon="Polygon" in annotation_mode,78        with_label="Mark" in annotation_mode)79    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)80    state = {81        ANNOTATED_IMAGE_KEY: annotated_image,82        DETECTIONS_KEY: detections83    }84    return (annotated_image, []), state85 86 87def prompt(88    message: str,89    history: List[List[str]],90    state: Dict[str, Any],91    api_key: Optional[str]92) -> str:93    if api_key == "":94        return "⚠️ Please set your OpenAI API key first"95    if state is None or ANNOTATED_IMAGE_KEY not in state:96        return "⚠️ Please generate SoM visual prompt first"97    return som.prompt_image(98        api_key=api_key,99        image=cv2.cvtColor(state[ANNOTATED_IMAGE_KEY], cv2.COLOR_BGR2RGB),100        prompt=message101    )102 103 104def on_image_input_clear():105    return None, {}106 107 108def highlight(109    state: Dict[str, Any],110    history: List[List[str]]111) -> Optional[Tuple[np.ndarray, List[Tuple[np.ndarray, str]]]]:112    if DETECTIONS_KEY not in state or ANNOTATED_IMAGE_KEY not in state:113        return None114 115    detections: sv.Detections = state[DETECTIONS_KEY]116    annotated_image: np.ndarray = state[ANNOTATED_IMAGE_KEY]117 118    if len(history) == 0:119        return None120 121    text = history[-1][-1]122    relevant_masks = som.extract_relevant_masks(123        text=text,124        detections=detections125    )126    relevant_masks = [127        (mask, mark)128        for mark, mask129        in relevant_masks.items()130    ]131    return annotated_image, relevant_masks132 133 134image_input = gr.Image(135    label="Input",136    type="numpy",137    tool="sketch",138    interactive=True,139    brush_radius=20.0,140    brush_color="#FFFFFF",141    height=512142)143checkbox_annotation_mode = gr.CheckboxGroup(144    choices=["Mark", "Polygon", "Mask", "Box"],145    value=['Mark'],146    label="Annotation Mode")147slider_mask_alpha = gr.Slider(148    minimum=0,149    maximum=1,150    value=0.05,151    label="Mask Alpha")152image_output = gr.AnnotatedImage(153    label="SoM Visual Prompt",154    color_map={155        str(i): sv.ColorPalette.default().by_idx(i).as_hex()156        for i in range(64)157    },158    height=512159)160openai_api_key = gr.Textbox(161    show_label=False,162    placeholder="Before you start chatting, set your OpenAI API key here",163    lines=1,164    type="password")165chatbot = gr.Chatbot(166    label="GPT-4V + SoM",167    height=256)168generate_button = gr.Button("Generate Marks")169highlight_button = gr.Button("Highlight Marks")170 171with gr.Blocks() as demo:172    gr.Markdown(MARKDOWN)173    inference_state = gr.State({})174    with gr.Row():175        with gr.Column():176            image_input.render()177            with gr.Accordion(178                    label="Detailed prompt settings (e.g., mark type)",179                    open=False):180                with gr.Row():181                    checkbox_annotation_mode.render()182                with gr.Row():183                    slider_mask_alpha.render()184        with gr.Column():185            image_output.render()186            generate_button.render()187            highlight_button.render()188    with gr.Row():189        openai_api_key.render()190    with gr.Row():191        gr.ChatInterface(192            chatbot=chatbot,193            fn=prompt,194            additional_inputs=[inference_state, openai_api_key])195 196    generate_button.click(197        fn=inference,198        inputs=[image_input, checkbox_annotation_mode, slider_mask_alpha],199        outputs=[image_output, inference_state])200    image_input.clear(201        fn=on_image_input_clear,202        outputs=[image_output, inference_state]203    )204    highlight_button.click(205        fn=highlight,206        inputs=[inference_state, chatbot],207        outputs=[image_output])208 209demo.queue().launch(debug=False, show_error=True)210