app.py

import spaces
import multiprocessing as mp
import numpy as np
from PIL import Image
import torch
try:
    import detectron2
except:
    import os
    os.system('pip install git+https://github.com/facebookresearch/detectron2.git')

from detectron2.config import get_cfg
from detectron2.projects.deeplab import add_deeplab_config
from detectron2.data.detection_utils import read_image
from mask_adapter import add_maskformer2_config, add_fcclip_config, add_mask_adapter_config
from mask_adapter.sam_maskadapter import SAMVisualizationDemo, SAMPointVisualizationDemo
import gradio as gr
import open_clip
from sam2.build_sam import build_sam2
from mask_adapter.modeling.meta_arch.mask_adapter_head import build_mask_adapter




def setup_cfg(config_file):
    cfg = get_cfg()
    add_deeplab_config(cfg)
    add_maskformer2_config(cfg)
    add_fcclip_config(cfg)
    add_mask_adapter_config(cfg)
    cfg.merge_from_file(config_file)
    cfg.freeze()
    return cfg

@spaces.GPU
@torch.no_grad()
@torch.autocast(device_type="cuda", dtype=torch.float32)
def inference_automatic(input_img, class_names):
    mp.set_start_method("spawn", force=True)
    config_file = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'
    cfg = setup_cfg(config_file)
    
    demo = SAMVisualizationDemo(cfg, 0.8, sam2_model, clip_model,mask_adapter)
    
    class_names = class_names.split(',')
    img = read_image(input_img, format="BGR")
    if len(class_names) == 1:
        class_names.append('others')
    txts = [f'a photo of {cls_name}' for cls_name in class_names]
    text = open_clip.tokenize(txts)


    text_features = clip_model.encode_text(text.cuda())
    text_features /= text_features.norm(dim=-1, keepdim=True)
        
    _, visualized_output = demo.run_on_image(img, class_names,text_features)

    return Image.fromarray(np.uint8(visualized_output.get_image())).convert('RGB')

@spaces.GPU
@torch.no_grad()
@torch.autocast(device_type="cuda", dtype=torch.float32)
def inference_point(input_img, evt: gr.SelectData,):


    mp.set_start_method("spawn", force=True)
    
    x, y = evt.index[0], evt.index[1]
    points = [[x, y]]
    print(f"Selected point: {points}")
    
    config_file = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'
    cfg = setup_cfg(config_file)
    
    demo = SAMPointVisualizationDemo(cfg, 0.8, sam2_model, clip_model,mask_adapter)
    
    img = read_image(input_img, format="BGR")
    

    _, visualized_output = demo.run_on_image_with_points(img, points)
    return visualized_output


sam2_model = None
clip_model = None
mask_adapter = None

def initialize_models(sam_path, adapter_pth, model_cfg, cfg):
    cfg = setup_cfg(cfg)
    global sam2_model, clip_model, mask_adapter

    if sam2_model is None:
        sam2_model = build_sam2(model_cfg, sam_path, device="cpu", apply_postprocessing=False)
        sam2_model = sam2_model.to("cuda")
        print("SAM2 model initialized.")
    
    if clip_model is None:
        clip_model, _, _ = open_clip.create_model_and_transforms("convnext_large_d_320", pretrained="laion2b_s29b_b131k_ft_soup")
        clip_model = clip_model.eval()
        clip_model = clip_model.to("cuda")
        print("CLIP model initialized.")
    
    if mask_adapter is None:
        mask_adapter = build_mask_adapter(cfg, "MASKAdapterHead").to("cuda")
        mask_adapter = mask_adapter.eval()
        adapter_state_dict = torch.load(adapter_pth)
        mask_adapter.load_state_dict(adapter_state_dict)
        print("Mask Adapter model initialized.")

model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
sam_path = './sam2.1_hiera_large.pt'
adapter_pth = './model_0279999_with_sem_new.pth'
cfg = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'

initialize_models(sam_path, adapter_pth, model_cfg, cfg)

# Examples for testing
examples = [
    ['./demo/images/000000001025.jpg', 'dog, beach, trees, sea, sky, snow, person, rocks, buildings, birds, beach umbrella, beach chair'],
    ['./demo/images/ADE_val_00000979.jpg', 'sky,sea,mountain,pier,beach,island,,landscape,horizon'],
    ['./demo/images/ADE_val_00001200.jpg', 'bridge, mountains, trees, water, sky, buildings, boats, animals, flowers, waterfalls, grasslands, rocks'],
]

output_labels = ['segmentation map']

title = '<center><h2>Mask-Adapter + Segment Anything-2</h2></center>'

description = """
<b>Mask-Adapter: The Devil is in the Masks for Open-Vocabulary Segmentation</b><br>
Mask-Adapter effectively extends to SAM or SAM-2 without additional training, achieving impressive results across multiple open-vocabulary segmentation benchmarks.<br>
<div style="display: flex; gap: 20px;">
    <a href="https://arxiv.org/abs/2406.20076">
        <img src="https://img.shields.io/badge/arXiv-Paper-red" alt="arXiv Paper">
    </a>
    <a href="https://github.com/hustvl/MaskAdapter">
        <img src="https://img.shields.io/badge/GitHub-Code-blue" alt="GitHub Code">
    </a>
</div>
"""

with gr.Blocks() as demo:
    gr.Markdown(title)  # Title
    gr.Markdown(description)  # Description

    with gr.Tabs():
        with gr.TabItem("Automatic Mode"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type='filepath', label="Input Image")
                    class_names = gr.Textbox(lines=1, placeholder=None, label='Class Names')
                with gr.Column():
                    output_image = gr.Image(type="pil", label='Segmentation Map')

                    # Buttons below segmentation map (now placed under segmentation map)
                    run_button = gr.Button("Run Automatic Segmentation")
                    run_button.click(inference_automatic, inputs=[input_image, class_names], outputs=output_image)
                    
                    clear_button = gr.Button("Clear")
                    clear_button.click(lambda: None, inputs=None, outputs=output_image)
            
            with gr.Row():
                gr.Examples(examples=examples, inputs=[input_image, class_names], outputs=output_image)

        with gr.TabItem("Point Mode"):
            with gr.Row():  
                with gr.Column(): 
                    def init_state():
                        return []
                    input_image = gr.Image(type='filepath', label="Upload Image", interactive=True)
                    points_input = gr.State(value=init_state()) 

                with gr.Column(): 
                    output_image_point = gr.Image(type="pil", label='Segmentation Map')

            input_image.select(inference_point, inputs=[input_image], outputs=output_image_point)

            clear_button_point = gr.Button("Clear Segmentation Map")
            clear_button_point.click(lambda: None, inputs=None, outputs=output_image_point)
        


            
    # Example images below buttons

    demo.launch()