app.py

from collections import namedtuple
import altair as alt
import pandas as pd
import gradio as gr

from PIL import Image

import os
import json
import argparse
import numpy as np
from tqdm import tqdm
from collections import OrderedDict

import torch
import torch.nn.functional as F
from torch.utils import data
import torchvision.transforms as transform
from torch.nn.parallel.scatter_gather import gather


import cv2
import math
import types
import functools
import torchvision.transforms as torch_transforms
import copy
import itertools
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.colors as mplc
import matplotlib.figure as mplfigure
import matplotlib.patches as mpatches


import sys
from unittest import result
sys.path.insert(0, "Mask2Former")
import tempfile
from pathlib import Path

sys.path.insert(0, "third_party/CLIP/")
os.system(f"pip3 install -Ue third_party/CLIP/")

os.system(f"pip install git+https://github.com/facebookresearch/detectron2.git")


# import some common detectron2 utilities
from detectron2.config import CfgNode as CN
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer, ColorMode
from detectron2.data import MetadataCatalog
from detectron2.utils.file_io import PathManager
from detectron2.utils.logger import setup_logger
from detectron2.projects.deeplab import add_deeplab_config
from detectron2.structures import Boxes, ImageList, Instances, BitMasks

# import Mask2Former project
from mask2former import add_mask_former_config

setup_logger()
logger = setup_logger(name="freeseg")


class Predictor(DefaultPredictor):

    def forward(self, original_image, labels=None):
        with torch.no_grad(): 
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

            inputs = {"image": image, "height": height, "width": width}
            predictions = self.model([inputs], labels)[0]
            return predictions


def create_predictor(task_names):
    cfg = get_cfg()
    add_deeplab_config(cfg)
    add_mask_former_config(cfg)
    cfg.merge_from_file("configs/coco-stuff-164k-156/mask2former_R101c_alltask_bs32_60k.yaml")
    cfg.MODEL.WEIGHTS = 'checkpoints/model_demo.pth'
    cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON = True
    cfg.MODEL.DEVICE = "cpu"
    cfg.INPUT.TASK_NAME = [task.lower() for task in task_names]
    predictor = Predictor(cfg)

    return predictor



"""
# FreeSeg Demo
"""


title = "FreeSeg"
description = """
<p style='text-align: center'> <a href='https://freeseg.github.io/' target='_blank'>Project Page</a> | <a href='https://arxiv.org/abs/2303.17225' target='_blank'>Paper</a> | <a href='https://github.com/bytedance/FreeSeg' target='_blank'>Code</a> </p>

Gradio demo for FreeSeg: Unified, Universal and Open-Vocabulary Image Segmentation. \n
You may click on of the examples or upload your own image. \n

"""  # noqa

article = """
<p style='text-align: center'><a href='https://arxiv.org/abs/2303.17225' target='_blank'>FreeSeg: Unified, Universal and Open-Vocabulary Image Segmentation</a> | <a href='https://github.com/bytedance/FreeSeg' target='_blank'>Github Repo</a></p>
"""  # noqa

examples = [
    [
        "examples/cat.jpg",
        "cat, grass, stone, other",
        ["Semantic segmentation"],
    ],
    [
        "examples/bus.jpg",
        "bus, person, road, building, tree, sky, other",
        ["Semantic segmentation", "Instance segmentation", "Panoptic segmentation"],
    ]
]





# labels = []
# for label in input_labels.split(","):
#     labels.append(label.strip())
# print(labels)

# tasks = []
# for task in input_tasks.split(","):
#     tasks.append(task.strip())


# coco_metadata = MetadataCatalog.get("coco_2017_val_panoptic_with_sem_seg")
# coco_metadata.stuff_classes[:len(labels)] = labels



def inference(image_path, labels, task_list):

    predictor = create_predictor(task_list)

    coco_metadata = MetadataCatalog.get("coco_2017_val_panoptic_with_sem_seg")
    coco_metadata.stuff_classes[:len(labels)] = labels

    image = Image.open(image_path)
    # image = np.array(image)
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    outputs = predictor.forward(image, labels)

    results = []
    if "Semantic segmentation" in task_list:
        sem_seg_out = outputs["sem_seg"].argmax(0).to("cpu")
        
        image_back = np.zeros_like(image)
        v = Visualizer(image[:, :, ::-1], coco_metadata, scale=0.6, instance_mode=ColorMode.IMAGE)
        semantic_result = v.draw_sem_seg(sem_seg_out, alpha=0.6).get_image()
        
        results.append(semantic_result)

    if "Panoptic segmentation" in task_list:
        coco_metadata.thing_classes[:len(labels)] = labels

        panvis = Visualizer(
                    image[:, :, ::-1],
                    coco_metadata,
                    scale=0.6, 
                    instance_mode=ColorMode.IMAGE
                )
        panoptic_seg, segments_info = outputs["panoptic_seg"]
        
        panvis_output = panvis.draw_panoptic_seg_predictions(
            panoptic_seg.cpu(), segments_info, alpha=0.6
        )
        
        panvis_output = panvis_output.get_image()
        results.append(panvis_output)
    
    if "Instance segmentation" in task_list:
        insvis = Visualizer(
            image[:, :, ::-1],
            coco_metadata,
            scale=0.6, 
            instance_mode=ColorMode.SEGMENTATION
        )
        instances = outputs["instances"].to(torch.device("cpu"))
        
        insvis_output = insvis.draw_instance_predictions(predictions=instances)

        insvis_output = insvis_output.get_image()
        results.insert(1, insvis_output)

    return results

    # logger.info("building class names")
    # demo_classes, demo_metadata = build_demo_classes_and_metadata(vocab, label_list)
    # with ExitStack() as stack:
    #     inference_model = OpenPanopticInference(
    #         model=model,
    #         labels=demo_classes,
    #         metadata=demo_metadata,
    #         semantic_on=False,
    #         instance_on=False,
    #         panoptic_on=True,
    #     )
    #     stack.enter_context(inference_context(inference_model))
    #     stack.enter_context(torch.no_grad())

    #     demo = VisualizationDemo(inference_model, demo_metadata, aug)
    #     img = utils.read_image(image_path, format="RGB")
    #     _, visualized_output = demo.run_on_image(img)
    #     return Image.fromarray(visualized_output.get_image())


# if uploaded_file is not None:
#     image = Image.open(uploaded_file)
#     # image = np.array(image)
#     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
#     outputs = predictor.forward(image, labels)


#     sem_seg_out = outputs["sem_seg"].argmax(0).to("cpu")
    
#     image_back = np.zeros_like(image)
#     v = Visualizer(image[:, :, ::-1], coco_metadata, scale=0.6, instance_mode=ColorMode.IMAGE)
#     semantic_result = v.draw_sem_seg(sem_seg_out, alpha=0.6).get_image()
    

#     coco_metadata.thing_classes[:len(labels)] = labels

#     panvis = Visualizer(
#                 image[:, :, ::-1],
#                 coco_metadata,
#                 scale=0.6, 
#                 instance_mode=ColorMode.IMAGE
#             )
#     panoptic_seg, segments_info = outputs["panoptic_seg"]
    
#     panvis_output = panvis.draw_panoptic_seg_predictions(
#         panoptic_seg.cpu(), segments_info, alpha=0.6
#     )
    
#     panvis_output = panvis_output.get_image()
    

#     insvis = Visualizer(
#         image[:, :, ::-1],
#         coco_metadata,
#         scale=0.6, 
#         instance_mode=ColorMode.SEGMENTATION
#     )
#     instances = outputs["instances"].to(torch.device("cpu"))
    
#     insvis_output = insvis.draw_instance_predictions(predictions=instances)

#     insvis_output = insvis_output.get_image()

#     back_result = np.ones_like(image) * 255
#     ######### show image online
#     image_show = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
#     fig = plt.figure()
#     plt.subplot(141)
#     plt.imshow(image_show)
#     plt.axis('off')
#     plt.title("Image", fontsize=6)
#     # plt.xlabel('Image')

#     plt.subplot(142)
#     task_name = tasks[0]
#     if task_name == "semantic segmentation":
#         plt.imshow(semantic_result)
#         # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#         plt.axis('off')
#         plt.title("Semantic", fontsize=6)
#     elif task_name == "instance segmentation":
#         plt.imshow(insvis_output)
#         # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#         plt.axis('off')
#         plt.title("Instance", fontsize=6)
#     elif task_name == "panoptic segmentation":
#         plt.imshow(panvis_output)
#         # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#         plt.axis('off')
#         plt.title("Panoptic", fontsize=6)
    
#     plt.subplot(143)
#     if len(tasks) > 1:
#         task_name = tasks[1]
#         if task_name == "semantic segmentation":
#             plt.imshow(semantic_result)
#             # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#             plt.axis('off')
#             plt.title("Semantic", fontsize=6)
#         elif task_name == "instance segmentation":
#             plt.imshow(insvis_output)
#             # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#             plt.axis('off')
#             plt.title("Instance", fontsize=6)
#         elif task_name == "panoptic segmentation":
#             plt.imshow(panvis_output)
#             # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#             plt.axis('off')
#             plt.title("Panoptic", fontsize=6)
#     else:
#         plt.axis('off')
#         plt.imshow(back_result)

#     plt.subplot(144)
#     if len(tasks) > 2:
#         task_name = tasks[2]
#         if task_name == "semantic segmentation":
#             plt.imshow(semantic_result)
#             # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#             plt.axis('off')
#             plt.title("Semantic", fontsize=6)
#         elif task_name == "instance segmentation":
#             plt.imshow(insvis_output)
#             # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#             plt.axis('off')
#             plt.title("Instance", fontsize=6)
#         elif task_name == "panoptic segmentation":
#             plt.imshow(panvis_output)
#             # plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
#             plt.axis('off')
#             plt.title("Panoptic", fontsize=6)
#     else:
#         plt.axis('off')
#         plt.imshow(back_result)

#     plt.tight_layout()

#     #st.image([image,seg], width=700, caption=["Input image", "Segmentation"])
#     st.pyplot(fig)
    

with gr.Blocks(title=title) as demo:
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>" + title + "</h1>")
    gr.Markdown(description)
    input_components = []
    output_components = []


    with gr.Row().style(equal_height=True, mobile_collapse=True):
        with gr.Column(scale=3, variant="panel") as input_component_column:
            input_image_gr = gr.inputs.Image(type="filepath")
            labels_gr = gr.inputs.Textbox(default="", label="Class labels")
            task_list_gr = gr.inputs.CheckboxGroup(
                choices=["Semantic segmentation", "Instance segmentation", "Panoptic segmentation"],
                default=["Semantic segmentation"],
                label="Task names",
            )
            input_components.extend([input_image_gr, labels_gr, task_list_gr])

    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
        clear_btn = gr.Button("Clear")
        

    with gr.Row():

        output_image_sem_gr = gr.outputs.Image(label="Semantic segmentation", type="pil")
        output_components.append(output_image_sem_gr)

        output_image_ins_gr = gr.outputs.Image(label="Instance segmentation", type="pil")
        output_components.append(output_image_ins_gr)

        output_image_pan_gr = gr.outputs.Image(label="Panoptic segmentation", type="pil")
        output_components.append(output_image_pan_gr)

    
    with gr.Column(scale=2):
        examples_handler = gr.Examples(
            examples=examples,
            inputs=[c for c in input_components if not isinstance(c, gr.State)],
            outputs=[c for c in output_components if not isinstance(c, gr.State)],
            fn=inference,
            cache_examples=torch.cuda.is_available(),
            examples_per_page=5,
        )

    gr.Markdown(article)

    submit_btn.click(
        inference,
        input_components,
        output_components,
        api_name="predict",
        scroll_to_output=True,
    )

    clear_btn.click(
        None,
        [],
        (input_components + output_components + [input_component_column]),
        _js=f"""() => {json.dumps(
                    [component.cleared_value if hasattr(component, "cleared_value") else None
                     for component in input_components + output_components] + (
                        [gr.Column.update(visible=True)]
                    )
                    + ([gr.Column.update(visible=False)])
                )}
                """,
    )

demo.launch()