OneFormer と OpenVINO によるユニバーサルなセグメント化#

import platform 

%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu "transformers>=4.26.0" "openvino>=2023.1.0" "nncf>=2.7.0" "gradio>=4.19" "torch>=2.1" scipy ipywidgets Pillow tqdm 

if platform.system() != "Windows":
     %pip install -q "matplotlib>=3.4" 
else:
     %pip install -q "matplotlib>=3.4,<3.7"

import warnings 
from collections import defaultdict 
from pathlib import Path 

from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation 
from transformers.models.oneformer.modeling_oneformer import ( 
    OneFormerForUniversalSegmentationOutput, 
) 
import torch 
import matplotlib.pyplot as plt 
import matplotlib.patches as mpatches 
from PIL import Image 
from PIL import ImageOps 

import openvino 

# `notebook_utils` モジュールを取得 
import requests 

r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py", 
) 

open("notebook_utils.py", "w").write(r.text) 
from notebook_utils import download_file

IR_PATH = Path("oneformer.xml") 
OUTPUT_NAMES = ["class_queries_logits", "masks_queries_logits"]

processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large") 
model = OneFormerForUniversalSegmentation.from_pretrained( 
    "shi-labs/oneformer_coco_swin_large", 
) 
id2label = model.config.id2label

task_seq_length = processor.task_seq_length 
shape = (800, 800) 
dummy_input = { 
    "pixel_values": torch.randn(1, 3, *shape), 
    "task_inputs": torch.randn(1, task_seq_length), 
}

model.config.torchscript = True 

if not IR_PATH.exists(): 
    with warnings.catch_warnings(): 
        warnings.simplefilter("ignore") 
        model = openvino.convert_model(model, example_input=dummy_input) 
    openvino.save_model(model, IR_PATH, compress_to_fp16=False)

import ipywidgets as widgets 

core = openvino.Core() 
device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

def prepare_inputs(image: Image.Image, task: str): 
    """Convert image to model input""" 
    image = ImageOps.pad(image, shape) 
    inputs = processor(image, [task], return_tensors="pt") 
    converted = { 
        "pixel_values": inputs["pixel_values"], 
        "task_inputs": inputs["task_inputs"], 
    } 
    return converted

def process_output(d): 
    """Convert OpenVINO model output to HuggingFace representation for visualization""" 
    hf_kwargs = {output_name: torch.tensor(d[output_name]) for output_name in OUTPUT_NAMES} 

    return OneFormerForUniversalSegmentationOutput(**hf_kwargs)

# ファイルからモデルを読み取り 
model = core.read_model(model=IR_PATH) 
# モデルをコンパイル 
compiled_model = core.compile_model(model=model, device_name=device.value)

class Visualizer:     @staticmethod 
    def extract_legend(handles): 
        fig = plt.figure() 
        fig.legend(handles=handles, ncol=len(handles) // 20 + 1, loc="center") 
        fig.tight_layout() 
        return fig 

    @staticmethod 
    def predicted_semantic_map_to_figure(predicted_map): 
        segmentation = predicted_map[0] 
        # 使用されているカラーマップを取得 
        viridis = plt.get_cmap("viridis", max(1, torch.max(segmentation))) 
        # すべての固有の番号を取得 
        labels_ids = torch.unique(segmentation).tolist() 
        fig, ax = plt.subplots() 
        ax.imshow(segmentation) 
        ax.set_axis_off() 
        handles = [] 
        for label_id in labels_ids: 
            label = id2label[label_id] 
            color = viridis(label_id) 
            handles.append(mpatches.Patch(color=color, label=label)) 
        fig_legend = Visualizer.extract_legend(handles=handles) 
        fig.tight_layout() 
        return fig, fig_legend 

    @staticmethod 
    def predicted_instance_map_to_figure(predicted_map): 
        segmentation = predicted_map[0]["segmentation"] 
        segments_info = predicted_map[0]["segments_info"] 
        # 使用されているカラーマップを取得 
        viridis = plt.get_cmap("viridis", max(torch.max(segmentation), 1)) 
        fig, ax = plt.subplots() 
        ax.imshow(segmentation) 
        ax.set_axis_off() 
        instances_counter = defaultdict(int) 
        handles = [] 
        # 各セグメントの凡例を描画 
        for segment in segments_info: 
            segment_id = segment["id"] 
            segment_label_id = segment["label_id"] 
            segment_label = id2label[segment_label_id] 
            label = f"{segment_label}-{instances_counter[segment_label_id]}" 
            instances_counter[segment_label_id] += 1 
            color = viridis(segment_id) 
            handles.append(mpatches.Patch(color=color, label=label)) 

        fig_legend = Visualizer.extract_legend(handles) 
        fig.tight_layout() 
        return fig, fig_legend 

    @staticmethod 
    def predicted_panoptic_map_to_figure(predicted_map): 
        segmentation = predicted_map[0]["segmentation"] 
        segments_info = predicted_map[0]["segments_info"] 
        # 使用されているカラーマップを取得 
        viridis = plt.get_cmap("viridis", max(torch.max(segmentation), 1)) 
        fig, ax = plt.subplots() 
        ax.imshow(segmentation) 
        ax.set_axis_off() 
        instances_counter = defaultdict(int) 
        handles = [] 
        # 各セグメントの凡例を描画 
        for segment in segments_info: 
            segment_id = segment["id"] 
            segment_label_id = segment["label_id"] 
            segment_label = id2label[segment_label_id] 
            label = f"{segment_label}-{instances_counter[segment_label_id]}" 
            instances_counter[segment_label_id] += 1 
            color = viridis(segment_id) 
            handles.append(mpatches.Patch(color=color, label=label)) 

        fig_legend = Visualizer.extract_legend(handles) 
        fig.tight_layout() 
        return fig, fig_legend 

    @staticmethod 
    def figures_to_images(fig, fig_legend, name_suffix=""): 
        seg_filename, leg_filename = ( 
            f"segmentation{name_suffix}.png", 
            f"legend{name_suffix}.png", 
        ) 
        fig.savefig(seg_filename, bbox_inches="tight") 
        fig_legend.savefig(leg_filename, bbox_inches="tight") 
        segmentation = Image.open(seg_filename) 
        legend = Image.open(leg_filename) 
        return segmentation, legend

def segment(model, img: Image.Image, task: str): 
    """ 
    Apply segmentation on an image.
    Args: 
        img: Input image. It will be resized to 800x800. 
        task: String describing the segmentation task. Supported values are: "semantic", "instance" and "panoptic". 
        Returns: 
            Tuple[Figure, Figure]: Segmentation map and legend charts.
    """ 
    if img is None: 
        raise gr.Error("Please load the image or use one from the examples list") 
    inputs = prepare_inputs(img, task) 
    outputs = model(inputs) 
    hf_output = process_output(outputs) 
    predicted_map = getattr(processor, f"post_process_{task}_segmentation")(hf_output, target_sizes=[img.size[::-1]]) 
    return getattr(Visualizer, f"predicted_{task}_map_to_figure")(predicted_map)

image = download_file("http://images.cocodataset.org/val2017/000000439180.jpg", "sample.jpg") 
image = Image.open("sample.jpg") 
image

from ipywidgets import Dropdown 

task = Dropdown(options=["semantic", "instance", "panoptic"], value="semantic") 
task

import matplotlib 

matplotlib.use("Agg") # 図の表示を無効にする 

def stack_images_horizontally(img1: Image, img2: Image): 
    res = Image.new("RGB", (img1.width + img2.width, max(img1.height, img2.height)), (255, 255, 255)) 
    res.paste(img1, (0, 0)) 
    res.paste(img2, (img1.width, 0)) 
    return res 

segmentation_fig, legend_fig = segment(compiled_model, image, task.value) 
segmentation_image, legend_image = Visualizer.figures_to_images(segmentation_fig, legend_fig) 
plt.close("all") 
prediction = stack_images_horizontally(segmentation_image, legend_image) 
prediction

compiled_quantized_model = None 

to_quantize = widgets.Checkbox( 
    value=False, 
    description="Quantization", 
    disabled=False, 
) 

to_quantize

# `skip_kernel_extension` モジュールを取得 
r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py", 
) 
open("skip_kernel_extension.py", "w").write(r.text) 

%load_ext skip_kernel_extension

import time 
import gradio as gr 

quantized_model_present = compiled_quantized_model is not None 

def compile_model(device): 
    global compiled_model 
    global compiled_quantized_model 
    compiled_model = core.compile_model(model=model, device_name=device) 
    if quantized_model_present: 
        compiled_quantized_model = core.compile_model(model=quantized_model, device_name=device) 

def segment_wrapper(image, task, run_quantized=False): 
    current_model = compiled_quantized_model if run_quantized else compiled_model 

    start_time = time.perf_counter() 
    segmentation_fig, legend_fig = segment(current_model, image, task) 
    end_time = time.perf_counter() 

    name_suffix = "" if not quantized_model_present else "_int8" if run_quantized else "_fp32" 
    segmentation_image, legend_image = Visualizer.figures_to_images(segmentation_fig, legend_fig, name_suffix=name_suffix) 
    plt.close("all") 
    result = stack_images_horizontally(segmentation_image, legend_image) 
    return result, f"{end_time - start_time:.2f}" 

with gr.Blocks() as demo: 
    with gr.Row(): 
        with gr.Column(): 
            inp_img = gr.Image(label="Image", type="pil") 
            inp_task = gr.Radio(["semantic", "instance", "panoptic"], label="Task", value="semantic") 
            inp_device = gr.Dropdown(label="Device", choices=core.available_devices + ["AUTO"], value="AUTO") 
        with gr.Column(): 
            out_result = gr.Image(label="Result (Original)" if quantized_model_present else "Result") 
            inference_time = gr.Textbox(label="Time (seconds)") 
            out_result_quantized = gr.Image(label="Result (Quantized)", visible=quantized_model_present) 
            inference_time_quantized = gr.Textbox(label="Time (seconds)", visible=quantized_model_present) 
    run_button = gr.Button(value="Run") 
    run_button.click( 
        segment_wrapper, 
        [inp_img, inp_task, gr.Number(0, visible=False)], 
        [out_result, inference_time], 
    ) 
    run_quantized_button = gr.Button(value="Run quantized", visible=quantized_model_present) 
    run_quantized_button.click( 
        segment_wrapper, 
        [inp_img, inp_task, gr.Number(1, visible=False)], 
        [out_result_quantized, inference_time_quantized], 
    ) 
    gr.Examples(examples=[["sample.jpg", "semantic"]], inputs=[inp_img, inp_task]) 

    def on_device_change_begin(): 
        return ( 
            run_button.update(value="Changing device...", interactive=False), 
            run_quantized_button.update(value="Changing device...", interactive=False), 
            inp_device.update(interactive=False), 
        ) 

    def on_device_change_end(): 
        return ( 
            run_button.update(value="Run", interactive=True), 
        run_quantized_button.update(value="Run quantized", interactive=True), 
        inp_device.update(interactive=True), 
        ) 

    inp_device.change(on_device_change_begin, outputs=[run_button, run_quantized_button, inp_device]).then(compile_model, inp_device).then( 
        on_device_change_end, outputs=[run_button, run_quantized_button, inp_device] 
    ) 

try: 
    demo.launch(debug=False) 
except Exception: 
    demo.launch(share=True, debug=False)
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細については、ドキュメントをご覧ください://gradio.app/docs/