CLIP と OpenVINO™ による言語と視覚の顕著性#

# 要件をインストール 
%pip install -q "openvino>=2023.1.0" 
%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu transformers "torch>=2.1" "gradio>=4.19"

from pathlib import Path 
from typing import Tuple, Union, Optional 
import requests 

from matplotlib import colors 
import matplotlib.pyplot as plt 
import numpy as np 
import torch 
import tqdm 
from PIL import Image 
from transformers import CLIPModel, CLIPProcessor

model_checkpoint = "openai/clip-vit-base-patch16" 

model = CLIPModel.from_pretrained(model_checkpoint).eval() 
processor = CLIPProcessor.from_pretrained(model_checkpoint)

def get_random_crop_params(image_height: int, image_width: int, min_crop_size: int) -> Tuple[int, int, int, int]: 
    crop_size = np.random.randint(min_crop_size, min(image_height, image_width)) 
    x = np.random.randint(image_width - crop_size + 1) 
    y = np.random.randint(image_height - crop_size + 1) 
    return x, y, crop_size 

def get_cropped_image(im_tensor: np.array, x: int, y: int, crop_size: int) -> np.array: 
    return im_tensor[y : y + crop_size, x : x + crop_size, ...] 

def update_saliency_map(saliency_map: np.array, similarity: float, x: int, y: int, crop_size: int) -> None: 
    saliency_map[ 
        y : y + crop_size, 
        x : x + crop_size,
     ] += similarity 

def cosine_similarity(one: Union[np.ndarray, torch.Tensor], other: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: 
    return one @ other.T / (np.linalg.norm(one) * np.linalg.norm(other))

n_iters = 300 
min_crop_size = 50 

query = "Who developed the Theory of General Relativity?" 
image_path = Path("example.jpg") 

r = requests.get("https://www.storypick.com/wp-content/uploads/2016/01/AE-2.jpg") 

with image_path.open("wb") as f: 
    f.write(r.content) 
image = Image.open(image_path) 
im_tensor = np.array(image) 

x_dim, y_dim = image.size

inputs = processor(text=[query], images=[im_tensor], return_tensors="pt") 
with torch.no_grad(): 
    results = model(**inputs) 
results.keys()

initial_similarity = cosine_similarity(results.text_embeds, results.image_embeds).item() # 1.  クエリーと画像の類似性の計算 
saliency_map = np.zeros((y_dim, x_dim)) 

for _ in tqdm.notebook.tqdm(range(n_iters)): # 6.  手順の繰り返し回数の設定 
    x, y, crop_size = get_random_crop_params(y_dim, x_dim, min_crop_size) 
    im_crop = get_cropped_image(im_tensor, x, y, crop_size) # 2. 画像をランダムに切り取る 

    inputs = processor(text=[query], images=[im_crop], return_tensors="pt") 
    with torch.no_grad(): 
        results = model(**inputs) # 3. 切り抜きとクエリーの類似性の計算 

    similarity = ( 
        cosine_similarity(results.text_embeds, results.image_embeds).item() - initial_similarity 
    ) # 4. 切り抜きとクエリーの類似性からクエリーと画像の類似性を減算 
    update_saliency_map(saliency_map, similarity, x, y, crop_size) # 5. 顕著性マップ上の領域の更新

plt.figure(dpi=150) 
plt.imshow(saliency_map, norm=colors.TwoSlopeNorm(vcenter=0), cmap="jet") 
plt.colorbar(location="bottom") 
plt.title(f'Query: "{query}"') 
plt.axis("off") 
plt.show()

def plot_saliency_map(image_tensor: np.ndarray, saliency_map: np.ndarray, query: Optional[str]) -> None: 
    fig = plt.figure(dpi=150) 
    plt.imshow(image_tensor) 
    plt.imshow( 
        saliency_map, 
        norm=colors.TwoSlopeNorm(vcenter=0), 
        cmap="jet", 
        alpha=0.5, # 顕著性マップを透明にして元の画像を表示 
    ) 
    if query: 
        plt.title(f'Query: "{query}"') 
    plt.axis("off") 
    return fig 

plot_saliency_map(im_tensor, saliency_map, query);

inputs = processor(text=[query], images=[im_tensor], return_tensors="pt") 
with torch.no_grad(): 
    results = model(**inputs) 
text_embeds = results.text_embeds # テキスト埋め込みを保存して後で使用 

initial_similarity = cosine_similarity(text_embeds, results.image_embeds).item() 
saliency_map = np.zeros((y_dim, x_dim)) 

for _ in tqdm.notebook.tqdm(range(n_iters)): 
    x, y, crop_size = get_random_crop_params(y_dim, x_dim, min_crop_size) 
    im_crop = get_cropped_image(im_tensor, x, y, crop_size) 

    image_inputs = processor(images=[im_crop], return_tensors="pt") # crop preprocessing 
    with torch.no_grad(): 
        image_embeds = model.get_image_features(**image_inputs) # 画像の埋め込みのみを計算 

    similarity = cosine_similarity(text_embeds, image_embeds).item() - initial_similarity 
    update_saliency_map(saliency_map, similarity, x, y, crop_size) 

plot_saliency_map(im_tensor, saliency_map, query);

import openvino as ov 

model_name = model_checkpoint.split("/")[-1] 

model.config.torchscript = True 
model.forward = model.get_text_features 
text_ov_model = ov.convert_model( 
    model, 
    example_input={ 
        "input_ids": inputs.input_ids, 
        "attention_mask": inputs.attention_mask, 
    }, 
) 

# プロセッサーから前処理後の画像サイズを取得 
crops_info = processor.image_processor.crop_size.values() if hasattr(processor, "image_processor") else 
processor.feature_extractor.crop_size.values() 
model.forward = model.get_image_features 
image_ov_model = ov.convert_model( 
    model, 
    example_input={"pixel_values": inputs.pixel_values}, 
    input=[1, 3, *crops_info], 
) 

ov_dir = Path("ir") 
ov_dir.mkdir(exist_ok=True) 
text_model_path = ov_dir / f"{model_name}_text.xml" 
image_model_path = ov_dir / f"{model_name}_image.xml" 

# 結果のモデルをディスクに書き込み 
ov.save_model(text_ov_model, text_model_path) 
ov.save_model(image_ov_model, image_model_path)

core = ov.Core() 

text_model = core.read_model(text_model_path) 
image_model = core.read_model(image_model_path)

from typing import Dict, Any 

image_model = core.read_model(image_model_path) 
    image_model = core.compile_model( model=image_model, 
    device_name=device.value, 
    config={"PERFORMANCE_HINT": "THROUGHPUT"}, 
)

text_inputs = dict(processor(text=[query], images=[im_tensor], return_tensors="np")) 
image_inputs = text_inputs.pop("pixel_values") 

text_embeds = text_model(text_inputs)[text_model.output()] 
image_embeds = image_model(image_inputs)[image_model.output()] 

initial_similarity = cosine_similarity(text_embeds, image_embeds) 
saliency_map = np.zeros((y_dim, x_dim))

def completion_callback( 
    infer_request: ov.InferRequest, # 推論結果 
    user_data: Dict[str, Any], # 入力ピクセル値とともに渡されたデータ 
) -> None: 
    pbar = user_data.pop("pbar") 

    image_embeds = infer_request.get_output_tensor().data 
    similarity = cosine_similarity(user_data.pop("text_embeds"), image_embeds) - user_data.pop("initial_similarity") 
    update_saliency_map(**user_data, similarity=similarity) 

    pbar.update(1) # 進捗バーを更新 

infer_queue = ov.AsyncInferQueue(image_model) 
infer_queue.set_callback(completion_callback)

def infer( 
    im_tensor, 
    x_dim, 
    y_dim, 
    text_embeds, 
    image_embeds, 
    initial_similarity, 
    saliency_map, 
    query, 
    n_iters, 
    min_crop_size, 
    _tqdm=tqdm.notebook.tqdm, 
    include_query=True, 
): 
    with _tqdm(total=n_iters) as pbar: 
        for _ in range(n_iters): 
            x, y, crop_size = get_random_crop_params(y_dim, x_dim, min_crop_size) 
            im_crop = get_cropped_image(im_tensor, x, y, crop_size) 

            image_inputs = processor(images=[im_crop], return_tensors="np") 

            # データをキューにプッシュ
             infer_queue.start_async( 
                # 推論データを通常通り渡す 
                image_inputs.pixel_values, 
                # 推論が完了した後、コールバックに渡されるデータ 
                { 
                    "text_embeds": text_embeds, 
                    "saliency_map": saliency_map, 
                    "initial_similarity": initial_similarity, 
                    "x": x, 
                    "y": y, 
                    "crop_size": crop_size, 
                    "pbar": pbar, 
                }, 
            ) 
        # すべてのデータをキューにプッシュした後、すべてのコールバックが完了するまで待機 
        infer_queue.wait_all() 

    return plot_saliency_map(im_tensor, saliency_map, query if include_query else None) 

infer( 
    im_tensor, 
    x_dim, 
    y_dim, 
    text_embeds, 
    image_embeds, 
    initial_similarity, 
    saliency_map, 
    query, 
    n_iters, 
    min_crop_size, 
    _tqdm=tqdm.notebook.tqdm, 
    include_query=True, 
);

import ipywidgets as widgets 

def build_saliency_map( 
    image: Image, 
    query: str, 
    n_iters: int = n_iters, 
    min_crop_size=min_crop_size, 
    _tqdm=tqdm.notebook.tqdm, 
    include_query=True, 
): 
    x_dim, y_dim = image.size 
    im_tensor = np.array(image) 

    text_inputs = dict(processor(text=[query], images=[im_tensor], return_tensors="np")) 
    image_inputs = text_inputs.pop("pixel_values") 

    text_embeds = text_model(text_inputs)[text_model.output()] 
    image_embeds = image_model(image_inputs)[image_model.output()] 

    initial_similarity = cosine_similarity(text_embeds, image_embeds) 
    saliency_map = np.zeros((y_dim, x_dim)) 

    return infer( 
        im_tensor, 
        x_dim, 
        y_dim, 
        text_embeds, 
        image_embeds, 
        initial_similarity, 
        saliency_map, 
        query, 
        n_iters, 
        min_crop_size, 
        _tqdm=_tqdm, 
        include_query=include_query, 
)

n_iters_widget = widgets.BoundedIntText( 
    value=n_iters, 
    min=1, 
    max=10000, 
    description="n_iters", 
) 
min_crop_size_widget = widgets.IntSlider( 
    value=min_crop_size, 
    min=1, 
    max=200, 
    description="min_crop_size", 
) 

@widgets.interact_manual(image_link="", query="", n_iters=n_iters_widget, min_crop_size=min_crop_size_widget) 
def build_saliency_map_from_image_link( 
    image_link: str, 
    query: str, 
    n_iters: int, 
    min_crop_size: int, 
) -> None: 
    try: 
        image_bytes = requests.get(image_link, stream=True).raw 
    except requests.RequestException as e: 
        print(f"Cannot load image from link: {image_link}\nException: {e}") 
        return 

    image = Image.open(image_bytes) 
    image = image.convert("RGB") # 透明チャネルを削除するか、グレースケール 1 チャネルを 3 チャネルに変換 

    build_saliency_map(image, query, n_iters, min_crop_size)

import io 

load_file_widget = widgets.FileUpload( 
    accept="image/*", 
    multiple=False, 
    description="Image file", 
) 

@widgets.interact_manual( 
    file=load_file_widget, 
    query="", 
    n_iters=n_iters_widget, 
    min_crop_size=min_crop_size_widget, 
) 
def build_saliency_map_from_file( 
    file: Path, 
    query: str = "", 
    n_iters: int = 2000, 
    min_crop_size: int = 50, 
) -> None: 
    image_bytes = io.BytesIO(file[0]["content"]) 
    try: 
        image = Image.open(image_bytes) 
    except Exception as e: 
        print(f"Cannot load the image: {e}") 
        return 

    image = image.convert("RGB") 

    build_saliency_map(image, query, n_iters, min_crop_size)

import gradio as gr 

def _process(image, query, n_iters, min_crop_size, _=gr.Progress(track_tqdm=True)): 
    saliency_map = build_saliency_map(image, query, n_iters, min_crop_size, _tqdm=tqdm.tqdm, include_query=False) 

    return saliency_map 

demo = gr.Interface( 
    _process, 
    [ 
        gr.Image(label="Image", type="pil"), 
        gr.Textbox(label="Query"), 
        gr.Slider(1, 10000, n_iters, label="Number of iterations"), 
        gr.Slider(1, 200, min_crop_size, label="Minimum crop size"), 
    ], 
    gr.Plot(label="Result"), 
    examples=[[image_path, query]], 
) 
try: 
    demo.queue().launch(debug=False) 
except Exception: 
    demo.queue().launch(share=True, debug=False) 
#リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/