| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import os |
| token = os.environ["GITHUB_TOKEN"] |
| os.system(f"pip install git+https://xvjiarui:{token}@github.com/xvjiarui/ODISE_NV.git") |
|
|
| import itertools |
| import json |
| from contextlib import ExitStack |
| import gradio as gr |
| import torch |
| from mask2former.data.datasets.register_ade20k_panoptic import ADE20K_150_CATEGORIES |
| from PIL import Image |
| from torch.cuda.amp import autocast |
|
|
| from detectron2.config import instantiate |
| from detectron2.data import MetadataCatalog |
| from detectron2.data import detection_utils as utils |
| from detectron2.data import transforms as T |
| from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES |
| from detectron2.evaluation import inference_context |
| from detectron2.utils.env import seed_all_rng |
| from detectron2.utils.logger import setup_logger |
| from detectron2.utils.visualizer import ColorMode, Visualizer, random_color |
|
|
| from odise import model_zoo |
| from odise.checkpoint import ODISECheckpointer |
| from odise.config import instantiate_odise |
| from odise.data import get_openseg_labels |
| from odise.modeling.wrapper import OpenPanopticInference |
| from odise.utils.file_io import ODISEHandler, PathManager |
| from odise.model_zoo.model_zoo import _ModelZooUrls |
|
|
| for k in ODISEHandler.URLS: |
| ODISEHandler.URLS[k] = ODISEHandler.URLS[k].replace("https://github.com/NVlabs/ODISE/releases/download/v1.0.0/", "https://huggingface.co/xvjiarui/download_cache/resolve/main/torch/odise/") |
| PathManager.register_handler(ODISEHandler()) |
| _ModelZooUrls.PREFIX = _ModelZooUrls.PREFIX.replace("https://github.com/NVlabs/ODISE/releases/download/v1.0.0/", "https://huggingface.co/xvjiarui/download_cache/resolve/main/torch/odise/") |
|
|
| setup_logger() |
| logger = setup_logger(name="odise") |
|
|
| COCO_THING_CLASSES = [ |
| label |
| for idx, label in enumerate(get_openseg_labels("coco_panoptic", True)) |
| if COCO_CATEGORIES[idx]["isthing"] == 1 |
| ] |
| COCO_THING_COLORS = [c["color"] for c in COCO_CATEGORIES if c["isthing"] == 1] |
| COCO_STUFF_CLASSES = [ |
| label |
| for idx, label in enumerate(get_openseg_labels("coco_panoptic", True)) |
| if COCO_CATEGORIES[idx]["isthing"] == 0 |
| ] |
| COCO_STUFF_COLORS = [c["color"] for c in COCO_CATEGORIES if c["isthing"] == 0] |
|
|
| ADE_THING_CLASSES = [ |
| label |
| for idx, label in enumerate(get_openseg_labels("ade20k_150", True)) |
| if ADE20K_150_CATEGORIES[idx]["isthing"] == 1 |
| ] |
| ADE_THING_COLORS = [c["color"] for c in ADE20K_150_CATEGORIES if c["isthing"] == 1] |
| ADE_STUFF_CLASSES = [ |
| label |
| for idx, label in enumerate(get_openseg_labels("ade20k_150", True)) |
| if ADE20K_150_CATEGORIES[idx]["isthing"] == 0 |
| ] |
| ADE_STUFF_COLORS = [c["color"] for c in ADE20K_150_CATEGORIES if c["isthing"] == 0] |
|
|
| LVIS_CLASSES = get_openseg_labels("lvis_1203", True) |
| |
| LVIS_COLORS = list( |
| itertools.islice(itertools.cycle([c["color"] for c in COCO_CATEGORIES]), len(LVIS_CLASSES)) |
| ) |
|
|
|
|
| class VisualizationDemo(object): |
| def __init__(self, model, metadata, aug, instance_mode=ColorMode.IMAGE): |
| """ |
| Args: |
| model (nn.Module): |
| metadata (MetadataCatalog): image metadata. |
| instance_mode (ColorMode): |
| parallel (bool): whether to run the model in different processes from visualization. |
| Useful since the visualization logic can be slow. |
| """ |
| self.model = model |
| self.metadata = metadata |
| self.aug = aug |
| self.cpu_device = torch.device("cpu") |
| self.instance_mode = instance_mode |
|
|
| def predict(self, original_image): |
| """ |
| Args: |
| original_image (np.ndarray): an image of shape (H, W, C) (in BGR order). |
| |
| Returns: |
| predictions (dict): |
| the output of the model for one image only. |
| See :doc:`/tutorials/models` for details about the format. |
| """ |
| height, width = original_image.shape[:2] |
| aug_input = T.AugInput(original_image, sem_seg=None) |
| self.aug(aug_input) |
| image = aug_input.image |
| image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) |
|
|
| inputs = {"image": image, "height": height, "width": width} |
| logger.info("forwarding") |
| with autocast(): |
| predictions = self.model([inputs])[0] |
| logger.info("done") |
| return predictions |
|
|
| def run_on_image(self, image): |
| """ |
| Args: |
| image (np.ndarray): an image of shape (H, W, C) (in BGR order). |
| This is the format used by OpenCV. |
| Returns: |
| predictions (dict): the output of the model. |
| vis_output (VisImage): the visualized image output. |
| """ |
| vis_output = None |
| predictions = self.predict(image) |
| visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode) |
| if "panoptic_seg" in predictions: |
| panoptic_seg, segments_info = predictions["panoptic_seg"] |
| vis_output = visualizer.draw_panoptic_seg( |
| panoptic_seg.to(self.cpu_device), segments_info |
| ) |
| else: |
| if "sem_seg" in predictions: |
| vis_output = visualizer.draw_sem_seg( |
| predictions["sem_seg"].argmax(dim=0).to(self.cpu_device) |
| ) |
| if "instances" in predictions: |
| instances = predictions["instances"].to(self.cpu_device) |
| vis_output = visualizer.draw_instance_predictions(predictions=instances) |
|
|
| return predictions, vis_output |
|
|
|
|
| cfg = model_zoo.get_config("Panoptic/odise_label_coco_50e.py", trained=True) |
|
|
| cfg.model.overlap_threshold = 0 |
| cfg.train.device = "cuda" if torch.cuda.is_available() else "cpu" |
| seed_all_rng(42) |
|
|
| dataset_cfg = cfg.dataloader.test |
| wrapper_cfg = cfg.dataloader.wrapper |
|
|
| aug = instantiate(dataset_cfg.mapper).augmentations |
|
|
| model = instantiate_odise(cfg.model) |
| model.to(torch.float16) |
| model.to(cfg.train.device) |
| ODISECheckpointer(model).load(cfg.train.init_checkpoint) |
|
|
|
|
| title = "ODISE" |
| description = """ |
| <p style='text-align: center'> <a href='https://jerryxu.net/ODISE' target='_blank'>Project Page</a> | <a href='https://arxiv.org/abs/2303.04803' target='_blank'>Paper</a> | <a href='https://github.com/NVlabs/ODISE' target='_blank'>Code</a> | <a href='https://youtu.be/Su7p5KYmcII' target='_blank'>Video</a></p> |
| |
| Gradio demo for ODISE: Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models. \n |
| You may click on of the examples or upload your own image. \n |
| |
| ODISE could perform open vocabulary segmentation, you may input more classes (separate by comma). |
| The expected format is 'a1,a2;b1,b2', where a1,a2 are synonyms vocabularies for the first class. |
| The first word will be displayed as the class name. |
| """ |
|
|
| article = """ |
| <p style='text-align: center'><a href='https://arxiv.org/abs/2303.04803' target='_blank'>Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models</a> | <a href='https://github.com/NVlab/ODISE' target='_blank'>Github Repo</a></p> |
| """ |
|
|
| examples = [ |
| [ |
| "demo/examples/coco.jpg", |
| "black pickup truck, pickup truck; blue sky, sky", |
| ["COCO (133 categories)", "ADE (150 categories)", "LVIS (1203 categories)"], |
| ], |
| [ |
| "demo/examples/ade.jpg", |
| "luggage, suitcase, baggage;handbag", |
| ["ADE (150 categories)"], |
| ], |
| [ |
| "demo/examples/ego4d.jpg", |
| "faucet, tap; kitchen paper, paper towels", |
| ["COCO (133 categories)"], |
| ], |
| ] |
|
|
|
|
| def build_demo_classes_and_metadata(vocab, label_list): |
| extra_classes = [] |
|
|
| if vocab: |
| for words in vocab.split(";"): |
| extra_classes.append([word.strip() for word in words.split(",")]) |
| extra_colors = [random_color(rgb=True, maximum=1) for _ in range(len(extra_classes))] |
|
|
| demo_thing_classes = extra_classes |
| demo_stuff_classes = [] |
| demo_thing_colors = extra_colors |
| demo_stuff_colors = [] |
|
|
| if any("COCO" in label for label in label_list): |
| demo_thing_classes += COCO_THING_CLASSES |
| demo_stuff_classes += COCO_STUFF_CLASSES |
| demo_thing_colors += COCO_THING_COLORS |
| demo_stuff_colors += COCO_STUFF_COLORS |
| if any("ADE" in label for label in label_list): |
| demo_thing_classes += ADE_THING_CLASSES |
| demo_stuff_classes += ADE_STUFF_CLASSES |
| demo_thing_colors += ADE_THING_COLORS |
| demo_stuff_colors += ADE_STUFF_COLORS |
| if any("LVIS" in label for label in label_list): |
| demo_thing_classes += LVIS_CLASSES |
| demo_thing_colors += LVIS_COLORS |
|
|
| MetadataCatalog.pop("odise_demo_metadata", None) |
| demo_metadata = MetadataCatalog.get("odise_demo_metadata") |
| demo_metadata.thing_classes = [c[0] for c in demo_thing_classes] |
| demo_metadata.stuff_classes = [ |
| *demo_metadata.thing_classes, |
| *[c[0] for c in demo_stuff_classes], |
| ] |
| demo_metadata.thing_colors = demo_thing_colors |
| demo_metadata.stuff_colors = demo_thing_colors + demo_stuff_colors |
| demo_metadata.stuff_dataset_id_to_contiguous_id = { |
| idx: idx for idx in range(len(demo_metadata.stuff_classes)) |
| } |
| demo_metadata.thing_dataset_id_to_contiguous_id = { |
| idx: idx for idx in range(len(demo_metadata.thing_classes)) |
| } |
|
|
| demo_classes = demo_thing_classes + demo_stuff_classes |
|
|
| return demo_classes, demo_metadata |
|
|
|
|
| def inference(image_path, vocab, label_list): |
|
|
| 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()) |
|
|
|
|
| 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(): |
| output_image_gr = gr.outputs.Image(label="Panoptic Segmentation", type="pil") |
| output_components.append(output_image_gr) |
|
|
| 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") |
| extra_vocab_gr = gr.inputs.Textbox(default="", label="Extra Vocabulary") |
| category_list_gr = gr.inputs.CheckboxGroup( |
| choices=["COCO (133 categories)", "ADE (150 categories)", "LVIS (1203 categories)"], |
| default=["COCO (133 categories)", "ADE (150 categories)", "LVIS (1203 categories)"], |
| label="Category to use", |
| ) |
| input_components.extend([input_image_gr, extra_vocab_gr, category_list_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, |
| ) |
| with gr.Row(): |
| clear_btn = gr.Button("Clear") |
| submit_btn = gr.Button("Submit", variant="primary") |
|
|
| 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() |
|
|
