MobileCLIP と OpenVINO を使用したビジュアルコンテンツ検索#

from pathlib import Path 

repo_dir = Path("./ml-mobileclip") 

if not repo_dir.exists():
     !git clone https://github.com/apple/ml-mobileclip.git

%pip install -q "./ml-mobileclip" --no-deps 

%pip install -q "clip-benchmark>=1.4.0" "datasets>=2.8.0" "open-clip-torch>=2.20.0" "timm>=0.9.5" "torch>=1.13.1" "torchvision>=0.14.1" --extra-index-url https://download.pytorch.org/whl/cpu 

%pip install -q "openvino>=2024.0.0" "gradio>=4.19" "matplotlib" "Pillow" "altair" "pandas" "opencv-python" "tqdm"

import ipywidgets as widgets 

model_dir = Path("checkpoints") 

supported_models = { 
    "MobileCLIP": { 
        "mobileclip_s0": { 
            "model_name": "mobileclip_s0", 
            "pretrained": model_dir / "mobileclip_s0.pt", 
            "url": "https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_s0.pt", 
            "image_size": 256, 
        }, 
        "mobileclip_s1": { 
            "model_name": "mobileclip_s1", 
            "pretrained": model_dir / "mobileclip_s1.pt", 
            "url": "https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_s1.pt", 
            "image_size": 256, 
        }, 
        "mobileclip_s2": { 
            "model_name": "mobileclip_s0", 
            "pretrained": model_dir / "mobileclip_s2.pt", 
            "url": "https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_s2.pt", 
            "image_size": 256, 
        }, 
        "mobileclip_b": { 
            "model_name": "mobileclip_b", 
            "pretrained": model_dir / "mobileclip_b.pt", 
            "url": "https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_b.pt", 
            "image_size": 224, 
        }, 
        "mobileclip_blt": { 
            "model_name": "mobileclip_b", 
            "pretrained": model_dir / "mobileclip_blt.pt", 
            "url": "https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_blt.pt", 
            "image_size": 224, 
        }, 
    }, 
    "CLIP":
            { "clip-vit-b-32": { 
            "model_name": "ViT-B-32", 
            "pretrained": "laion2b_s34b_b79k", 
            "image_size": 224, 
        }, 
        "clip-vit-b-16": { 
            "image_name": "ViT-B-16", 
            "pretrained": "openai", 
            "image_size": 224, 
        }, 
        "clip-vit-l-14": { 
            "image_name": "ViT-L-14", 
            "pretrained": "datacomp_xl_s13b_b90k", 
            "image_size": 224, 
        }, 
        "clip-vit-h-14": { 
            "image_name": "ViT-H-14", 
            "pretrained": "laion2b_s32b_b79k", 
            "image_size": 224, 
        }, 
    }, 
    "SigLIP": { 
        "siglip-vit-b-16": { 
            "model_name": "ViT-B-16-SigLIP", 
            "pretrained": "webli", 
            "image_size": 224, 
        }, 
        "siglip-vit-l-16": { 
            "model_name": "ViT-L-16-SigLIP-256", 
            "pretrained": "webli", 
            "image_size": 256, 
        }, 
    }, 
} 

model_type = widgets.Dropdown(options=supported_models.keys(), default="MobileCLIP", description="Model type:") 
model_type

available_models = supported_models[model_type.value] 

model_checkpoint = widgets.Dropdown( 
    options=available_models.keys(), 
    default=list(available_models), 
    description="Model:", 
) 

model_checkpoint

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 

model_config = available_models[model_checkpoint.value]

import types 
import torch.nn.functional as F 

def se_block_forward(self, inputs): 
    """Apply forward pass.""" 
    b, c, h, w = inputs.size() 
    x = F.avg_pool2d(inputs, kernel_size=[8, 8]) 
    x = self.reduce(x) 
    x = F.relu(x) 
    x = self.expand(x) 
    x = torch.sigmoid(x) 
    x = x.view(-1, c, 1, 1) 
    return inputs * x

import openvino as ov 
import gc 

ov_models_dir = Path("ov_models") 
ov_models_dir.mkdir(exist_ok=True) 

image_encoder_path = ov_models_dir / f"{model_checkpoint.value}_im_encoder.xml" 

if not image_encoder_path.exists(): 
    if "mobileclip_s" in model_name: 
        model.image_encoder.model.conv_exp.se.forward = types.MethodType(se_block_forward, model.image_encoder.model.conv_exp.se) 
    model.forward = model.encode_image 
    ov_image_encoder = ov.convert_model( 
        model, 
        example_input=image_tensor, 
        input=[-1, 3, image_tensor.shape[2], image_tensor.shape[3]], 
    ) 
    ov.save_model(ov_image_encoder, image_encoder_path) 
    del ov_image_encoder 
    gc.collect() 

text_encoder_path = ov_models_dir / f"{model_checkpoint.value}_text_encoder.xml" 

if not text_encoder_path.exists(): 
    model.forward = model.encode_text 
    ov_text_encoder = ov.convert_model(model, example_input=text, input=[-1, text.shape[1]]) 
    ov.save_model(ov_text_encoder, text_encoder_path) 
    del ov_text_encoder 
    gc.collect() 

del model 
gc.collect();

import altair as alt 
import cv2 
import gradio as gr 
import pandas as pd 
import torch 
from PIL import Image 
from torch.utils.data import DataLoader, Dataset 
from torchvision.transforms.functional import to_pil_image, to_tensor 
from torchvision.transforms import ( 
    CenterCrop, 
    Compose, 
    InterpolationMode, 
    Resize, 
    ToTensor, 
) 
from open_clip.transform import image_transform 

current_device = device.value 
current_model = image_encoder_path.name.split("_im_encoder")[0] 

available_converted_models = [model_file.name.split("_im_encoder")[0] for model_file in 
ov_models_dir.glob("*_im_encoder.xml")] 
available_devices = list(core.available_devices) + ["AUTO"] 

download_file( 
    "https://github.com/intel-iot-devkit/sample-videos/raw/master/car-detection.mp4", 
    directory=sample_path, 
) 
download_file( 

"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Coco%20Walking%20in%20Berkeley.mp4", 
    directory=sample_path, 
    filename="coco.mp4", 
) 

def get_preprocess_and_tokenizer(model_name): 
    if "mobileclip" in model_name: 
        resolution = supported_models["MobileCLIP"][model_name]["image_size"] 
        resize_size = resolution 
        centercrop_size = resolution 
        aug_list = [ 
            Resize( 
                resize_size, 
                interpolation=InterpolationMode.BILINEAR, 
            ), 
            CenterCrop(centercrop_size), 
            ToTensor(), 
        ] 
        preprocess = Compose(aug_list) 
        tokenizer = mobileclip.get_tokenizer(supported_models["MobileCLIP"][model_name]["model_name"]) 
    else: 
        model_configs = supported_models["SigLIP"] if "siglip" in model_name else supported_models["CLIP"] 
        resize_size = model_configs[model_name]["image_size"] 
        preprocess = image_transform((resize_size, resize_size), is_train=False, resize_mode="longest") 
        tokenizer = open_clip.get_tokenizer(model_configs[model_name]["model_name"]) 
    return preprocess, tokenizer 

def run( 
    path: str, 
    text_search: str, 
    image_search: Image.Image, 
    model_name: str, 
    device: str, 
    thresh: float, 
    stride: int, 
    batch_size: int, 
): 
    assert path, "An input video should be provided" 
    assert text_search is not None or image_search is not None, "A text or image query should be provided" 
    global current_model 
    global current_device 
    global preprocess 
    global tokenizer 
    global ov_compiled_image_encoder 
    global ov_compiled_text_encoder 

    if current_model != model_name or device != current_device: 
        ov_compiled_image_encoder = core.compile_model(ov_models_dir / f"{model_name}_im_encoder.xml", device) 
        ov_compiled_text_encoder = core.compile_model(ov_models_dir / f"{model_name}_text_encoder.xml", device) 
        preprocess, tokenizer = get_preprocess_and_tokenizer(model_name) 
        current_model = model_name 
        current_device = device 
    # ビデオをロード 
    dataset = LoadVideo(path, transforms=preprocess, vid_stride=stride) 
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0) 

    # 画像クエリー機能を取得 
    if image_search: 
        image = preprocess(image_search).unsqueeze(0) 
        query_features = torch.from_numpy(ov_compiled_image_encoder(image)[0]) 
        query_features /= query_features.norm(dim=-1, keepdim=True) 
    # テキストクエリー機能を取得 
    else:
        # 検索フレーズをトークン化 
        text = tokenizer([text_search]) 
        # テキストクエリーをエンコード 
        query_features = torch.from_numpy(ov_compiled_text_encoder(text)[0]) 
        query_features /= query_features.norm(dim=-1, keepdim=True) 
    # 各フレームをエンコードし、クエリー機能と比較 
    matches = [] 
    matches_probs = [] 
    res = pd.DataFrame(columns=["Frame", "Timestamp", "Similarity"]) 
    for image, orig, frame, timestamp in dataloader: 
        with torch.no_grad(): 
            image_features = torch.from_numpy(ov_compiled_image_encoder(image)[0]) 

        image_features /= image_features.norm(dim=-1, keepdim=True) 
        probs = query_features.cpu().numpy() @ image_features.cpu().numpy().T 
        probs = probs[0] 

        # フレーム類似性値を保存 
        df = pd.DataFrame( 
            { 
                "Frame": frame.tolist(), 
                "Timestamp": torch.round(timestamp / 1000, decimals=2).tolist(), 
                "Similarity": probs.tolist(), 
            } 
        ) 
        res = pd.concat([res, df]) 

        # フレームが閾値を超えていないか確認 
        for i, p in enumerate(probs): 
            if p > thresh: 
                matches.append(to_pil_image(orig[i])) 
                matches_probs.append(p) 

        print(f"Frames: {frame.tolist()} - Probs: {probs}") 

# 類似性値のプロットを作成 
lines = ( 
    alt.Chart(res) 
    .mark_line(color="firebrick") 
    .encode( 
        alt.X("Timestamp", title="Timestamp (seconds)"), 
        alt.Y("Similarity", scale=alt.Scale(zero=False)), 
    ) 
).properties(width=600) 
rule = alt.Chart().mark_rule(strokeDash=[6, 3], size=2).encode(y=alt.datum(thresh)) 

selected_frames = np.argsort(-1 * np.array(matches_probs))[:20] 
matched_sorted_frames = [matches[idx] for idx in selected_frames] 

return ( 
    lines + rule, 
    matched_sorted_frames, 
) # UI がクラッシュしないよう、最大​​ 20 枚の画像のみを返す 

class LoadVideo(Dataset): 
    def __init__(self, path, transforms, vid_stride=1): 
        self.transforms = transforms 
        self.vid_stride = vid_stride 
        self.cur_frame = 0 
        self.cap = cv2.VideoCapture(path) 
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride) 

def __getitem__(self, _):
    # ビデオを読み込み 
    # フレームをスキップ 
    for _ in range(self.vid_stride): 
        self.cap.grab() 
        self.cur_frame += 1 

    # フレームを読み込み 
    _, img = self.cap.retrieve() 
    timestamp = self.cap.get(cv2.CAP_PROP_POS_MSEC) 

    # PIL に変換 
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 
    img = Image.fromarray(np.uint8(img)) 

    # トランスフォームを適用 
    img_t = self.transforms(img) 

    return img_t, to_tensor(img), self.cur_frame, timestamp 

def __len__(self): 
    return self.total_frames 

desc_text = """ 
Search the content's of a video with a text description. __Note__: Long videos (over a few minutes) may cause UI performance issues.
    """ 
text_app = gr.Interface( 
    description=desc_text, 
    fn=run, 
    inputs=[ 
        gr.Video(label="Video"), 
        gr.Textbox(label="Text Search Query"), 
        gr.Image(label="Image Search Query", visible=False), 
        gr.Dropdown( 
            label="Model", 
            choices=available_converted_models, 
            value=model_checkpoint.value, 
        ), 
        gr.Dropdown(label="Device", choices=available_devices, value=device.value), 
        gr.Slider(label="Threshold", maximum=1.0, value=0.2), 
        gr.Slider(label="Frame-rate Stride", value=4, step=1), 
        gr.Slider(label="Batch Size", value=4, step=1), 
    ], 
    outputs=[ 
        gr.Plot(label="Similarity Plot"), 
        gr.Gallery(label="Matched Frames", columns=2, object_fit="contain", height="auto"), 
    ], 
    examples=[[sample_path / "car-detection.mp4", "white car"]], 
    allow_flagging="never", 
) 

desc_image = """ 
Search the content's of a video with an image query.__Note__: Long videos (over a few minutes) may cause UI performance issues.
    """ 
image_app = gr.Interface( 
    description=desc_image, 
    fn=run, 
    inputs=[ 
        gr.Video(label="Video"), 
        gr.Textbox(label="Text Search Query", visible=False), 
        gr.Image(label="Image Search Query", type="pil"), 
        gr.Dropdown( 
            label="Model", 
            choices=available_converted_models, 
            value=model_checkpoint.value, 
        ), 
        gr.Dropdown(label="Device", choices=available_devices, value=device.value), 
        gr.Slider(label="Threshold", maximum=1.0, value=0.2), 
        gr.Slider(label="Frame-rate Stride", value=4, step=1), 
        gr.Slider(label="Batch Size", value=4, step=1), 
    ], 
    outputs=[ 
        gr.Plot(label="Similarity Plot"), 
        gr.Gallery(label="Matched Frames", columns=2, object_fit="contain", height="auto"), 
    ], 
    allow_flagging="never", 
    examples=[[sample_path / "coco.mp4", None, sample_path / "dog.png"]], 
) 
demo = gr.TabbedInterface( 
    interface_list=[text_app, image_app], 
    tab_names=["Text Query Search", "Image Query Search"], 
    title="CLIP Video Content Search", 
) 

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') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/