なんでもボックス検出できるowlvit

すぐ使えてなんでも検出

またすごいの出てきたよ。
物の名前を入れたらなんでも検出してくれる。

たとえば、「サーフボード（青ボックス）」と「波（赤ボックス）」を検出してみる。

あるいは、「横笛」を検出してみる。

まあたいしたものである。

お気づきかもしれないが、COCO datasetにないクラスも余裕で探せる。

しかもCPUでも２、３秒。
GPUで実行すれば0.5秒以下である（T4）。
クラス数が多くなるとだんだん処理も長くなるけど。

使い方

モデルの初期化

from PIL import Image, ImageDraw, ImageFont
import torch
from transformers import OwlViTProcessor, OwlViTForObjectDetection

# モデルとプロセッサの読み込み
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
processor.cuda()
model.cuda()

実行

import time
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import OwlViTProcessor, OwlViTForObjectDetection

# 画像の読み込み
image = Image.open("animal.jpg")
texts = [["a photo of a dog",
          "a photo of a cat"]]

# モデル処理
start = time.time()
inputs = processor(text=texts, images=image, return_tensors="pt")
inputs = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
outputs = model(**inputs)
print(f"Inference time: {time.time() - start} seconds")

# ボックスのスケーリング設定
target_sizes = torch.Tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes, threshold=0.1)

# ラベルごとの色設定
label_colors = {
    "a photo of a dog": "red",
    "a photo of a cat": "blue"
}

# 描画設定
draw = ImageDraw.Draw(image)
font = ImageFont.load_default()  # デフォルトフォントを使用

# 各検出結果にボックスとラベルを描画
i = 0  # 1枚目の画像
text = texts[i]
boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"]

for box, score, label in zip(boxes, scores, labels):
    if score >= 0.2:  # 信頼度0.2以上のみ描画
        box = [round(i, 2) for i in box.tolist()]
        label_text = text[label]
        color = label_colors.get(label_text, "red")  # デフォルト色を赤に設定

        # ボックスとラベルの描画
        draw.rectangle(box, outline=color, width=3)  # ボックスの描画
        draw.text((box[0], box[1] - 10), f"{label_text}: {round(score.item(), 2)}", fill=color, font=font)  # ラベルの描画

# 結果の保存
image.save("output.jpg")  # 画像を保存

みんな仕事いそがしいと思うけど、
てきとうに仕事以外のこともやって、
人生がんばっていこうね。

🐣

---

フリーランスエンジニアです。
お仕事のご相談こちらまで
[e-mail]

Core MLを使ったアプリを作っています。
機械学習関連の情報を発信しています。

Twitter
Medium