この動画で使っているプログラムを、こちらの記事に転記しています。

記事の概要

こちらの記事でGoogle Colab用に実装されている、DETRのファインチューニングプログラムをWindows用にカスタマイズしています。

また、自前で用意した画像データに対して、ラベリング（labelme）・ラベルの確認（coco-viewer）用の手順を追加し、自前データを用いたファインチューニングの手順も、別の記事として作成しています。

環境

OS:Windows 11
GPU:GeForce RTX 4090
CPU:i9-13900KF
memory:64G
python:3.10.10
pytorch:2.0.1
CUDA:11.8
cuDNN:8.8

以下の環境でも動作確認済み。
GPU:GeForce RTX 3060 laptop
CPU:i7-10750H
memory:16G

参考記事のWindows版

以下のコマンドを実行し、環境構築を構築する。

# 仮想環境を構築・アクティベイト
python -m venv detr_env
cd detr_env\Scripts
activate
cd ..

# 対象のライブラリをインストール
git clone https://github.com/EscVM/OIDv4_ToolKit.git
pip install urllib3==1.25.11 folium==0.2.1
pip install -r OIDv4_ToolKit/requirements.txt

# 学習用データをダウンロード
python OIDv4_ToolKit/main.py downloader -y --classes Apple Orange Ball Balloon Clock --type_csv train
python OIDv4_ToolKit/main.py downloader -y --classes Apple Orange Ball Balloon Clock --type_csv validation

以下のプログラムを実行し、学習用データをCOCO formatに変換する。

import os, json, glob

def OID2JSON(OIDFiles, saveName, subset):
    """
    アノテーションをOpenImage format(txt)からCOCO format(json)に変換
    Parameters
    ----------
    OIDFiles : string
        OpenImageDatasetのフォルダパス
    saveName : string
        保存ファイル名(json)
    subset : string
        変換したいtype_csv。train、validation、testのいずれかを指定。
    """
    attrDict = dict()
    # categories要素の設定
    attrDict['categories'] = []
    categories = sorted(os.listdir(os.path.join(OIDFiles, 'Dataset', subset)))
    for i in range(len(categories)):
        attrDict['categories'].append({'supercategory': 'none', 'id': i, 'name': categories[i]})
    
    images = list()
    annotations = list()
    filenames = list()
    image_id = 1
    anno_id = 1
    for category in attrDict['categories']:
        for jpg_file in glob.glob(os.path.join(OIDFiles, 'Dataset', subset, category['name'], '*.jpg')):
            filename = os.path.splitext(os.path.basename(jpg_file))[0]
            # カテゴリ全体で同じファイル名が存在する場合、imageとannoをリネーム
            if filename in filenames:
                rename_filename = filename + '_' + str(image_id)
                os.rename(jpg_file, os.path.join(OIDFiles, 'Dataset', subset, category['name'], rename_filename + '.jpg'))
                os.rename(os.path.join(OIDFiles, 'Dataset', subset, category['name'], 'Label', filename + '.txt'),
                                    os.path.join(OIDFiles, 'Dataset', subset, category['name'], 'Label', rename_filename + '.txt'))
                filename = rename_filename
            
            filenames.append(filename)
            # images要素の設定
            # ※DETRではheightとwidthを使わないので、'none'を設定
            image = {'file_name': filename + '.jpg', 'height': 'none', 'width': 'none', 'id': image_id}
            images.append(image)
            # annotations要素の設定
            anno_path = os.path.join(OIDFiles, 'Dataset', subset, category['name'], 'Label', filename + '.txt')
            with open(anno_path) as f:
                for line in f:
                    splitline = line.split(' ')
                    # カテゴリがcategories要素に存在しないバウンディングボックスは使わない
                    if splitline[0] in [d.get('name') for d in attrDict['categories']]:
                        # OpenImageの座標は(xmin, ymin, xmax, ymax)、COCOの座標は(x, y, width, height)
                        x1 = int(float(splitline[1]))
                        y1 = int(float(splitline[2]))
                        x2 = int(float(splitline[3])) - x1
                        y2 = int(float(splitline[4])) - y1
                        # areaはピクセル数(float)
                        area = float(x2 * y2)
                        # segmentationは(x1, y1, x2, y2, ...)と順番に定義
                        segmentation = [[x1, y1, x1, (y1+y2), (x1+x2), (y1+y2), (x1+x2), y1]]
                        annotation = {'iscrowd': 0, 'image_id': image_id, 'bbox': [x1, y1, x2, y2], 'area': area,
                                      'category_id': category['id'], 'ignore': 0, 'id': anno_id, 'segmentation': segmentation}
                        anno_id += 1
                        annotations.append(annotation)
                
            image_id = image_id + 1
    
    attrDict['images'] = images
    attrDict['annotations'] = annotations
    attrDict['type'] = 'instances'
    jsonString = json.dumps(attrDict)
    with open(saveName, 'w') as f:
        f.write(jsonString)


OID2JSON('OID', 'custom_train.json', 'train')
OID2JSON('OID', 'custom_val.json', 'validation')

以下のプログラムを実行し、学習用データを移動する。

import shutil
import glob
import os

source_train_paths = glob.glob(os.path.join('OID\\Dataset', 'train', '**\\'))
source_val_paths = glob.glob(os.path.join('OID\\Dataset', 'validation', '**\\'))
train_path = 'data\\custom\\train2017\\'
val_path = 'data\\custom\\val2017\\'
convert_anno_path = 'data\\custom\\annotations\\'
# ディレクトリ作成
os.makedirs(train_path, exist_ok=True)
os.makedirs(val_path, exist_ok=True)
os.makedirs(convert_anno_path, exist_ok=True)
# train移動
for source_train_path in source_train_paths:
    for img_path in glob.glob(os.path.join(source_train_path, '*.jpg')):
        shutil.move(img_path, train_path)

# val移動
for source_val_path in source_val_paths:
    for img_path in glob.glob(os.path.join(source_val_path, '*.jpg')):
        shutil.move(img_path, val_path)

# anno移動
shutil.move('custom_train.json', convert_anno_path)
shutil.move('custom_val.json', convert_anno_path)

以下のコマンドで必要なライブラリをインストールする。

pip install torch torchvision torchtext -f https://download.pytorch.org/whl/cu118/torch_stable.html
pip install torch torchvision torchtext
pip install matplotlib
pip install pycocotools
pip install scipy

以下のコマンドを実行し、DETRのライブラリを取得する。

# detr_envフォルダ配下で
rd /s detr
git clone https://github.com/woctezuma/detr.git

# ブランチの切替
cd detr/
git checkout finetune

以下のプログラムを実行し、学習済みモデルを保存する。

import torch, torchvision
import torchvision.transforms as T
import matplotlib.pyplot as plt
from PIL import Image
import requests

# 学習済みモデルの取得
checkpoint = torch.hub.load_state_dict_from_url(
    url='https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth',
    map_location='cpu',
    check_hash=True
)

# 分類ヘッドの削除
del checkpoint['model']['class_embed.weight']
del checkpoint['model']['class_embed.bias']

# 保存
torch.save(checkpoint, 'detr-r50_no-class-head.pth')

以下のコマンドで出力先フォルダを作成する。

rd /s outputs
mkdir outputs

以下のコマンドでチューニングを実施する。（制度を向上するなら、epochsを50~100などに変更してもよい。）

python main.py --dataset_file "custom" --coco_path "..\\data\\custom\\" --output_dir "outputs" --resume "detr-r50_no-class-head.pth" --num_classes 5 --epochs 15

以下のプログラムを実行し、学習結果を確認する。

import torch, torchvision
import torchvision.transforms as T
import matplotlib.pyplot as plt
from pathlib import Path
from io import BytesIO
from PIL import Image
import requests

# %cd /content/detr/
# log_directory = [Path('/content/detr/outputs')]
log_directory = [Path('\\outputs')]

# 実線 ... トレーニング結果(train_loss)
# 破線 ... 検証結果(val_loss)
fields_of_interest = (
    'loss',
    'mAP',
)
# plot_logs(log_directory, fields_of_interest)

finetuned_model = torch.hub.load('facebookresearch/detr',
                       'detr_resnet50',
                       pretrained=False,
                       num_classes=5)
checkpoint = torch.load('detr\\outputs\\checkpoint_epoch.pth',
                        map_location='cpu')
finetuned_model.load_state_dict(checkpoint['model'], strict=False)
finetuned_model.eval()

original_model = torch.hub.load('facebookresearch/detr', 'detr_resnet50_dc5', pretrained=True)
original_model.eval()

# 可視化用クラスラベル
oid_labels = [
  'Apple',
  'Ball',
  'Balloon',
  'Clock',
  'Orange',
]
coco_labels = [
    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'
]
# 可視化用COLOR
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]

# 標準的なPyTorchのmean-std入力画像の正規化
transform = T.Compose([
    T.Resize(800),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def box_cxcywh_to_xyxy(x):
    """
    (center_x, center_y, width, height)から(xmin, ymin, xmax, ymax)に座標変換
    """
    # unbind(1)でTensor次元を削除
    # (center_x, center_y, width, height)*N → (center_x*N, center_y*N, width*N, height*N)
    x_c, y_c, w, h = x.unbind(1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
    # (center_x, center_y, width, height)*N の形に戻す
    return torch.stack(b, dim=1)

def rescale_bboxes(out_bbox, size):
    """
    バウンディングボックスのリスケール
    """
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    # バウンディングボックスの[0～1]から元画像の大きさにリスケール
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
    return b

def filter_bboxes_from_outputs(outputs, threshold=0.7):
    # 閾値以上の信頼度を持つ予測値のみを保持
    probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > threshold
    probas_to_keep = probas[keep]
    # [0, 1]のボックスを画像のスケールに変換
    bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
    return probas_to_keep, bboxes_scaled

# 結果の表示
def plot_finetuned_results(pil_img, prob=None, boxes=None, labels=None):
    plt.figure(figsize=(16, 10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    if prob is not None and boxes is not None:
        for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
            ax.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin,
                                                                 fill=False, color=c, linewidth=3))
            cl = p.argmax()
            print(labels, p)
            text = f'{labels[cl]}: {p[cl]:0.2f}'
            ax.text(xmin, ymin, text, fontsize=15,
                            bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()

# 物体検出
def run_worflow(my_image, my_model, labels, threshold=0.7):
    # mean-std入力画像の正規化(バッチサイズ : 1)
    img = transform(my_image).unsqueeze(0)
    
    # モデルに反映
    outputs = my_model(img)
    probas_to_keep, bboxes_scaled = filter_bboxes_from_outputs(outputs, threshold=threshold)
    plot_finetuned_results(my_image, probas_to_keep, bboxes_scaled, labels)

url = 'https://farm7.staticflickr.com/52/106887535_a29c34113b_o.jpg'
response = requests.get(url)
im = Image.open(BytesIO(response.content))

# Fine-Tuningモデルで物体検出(閾値0.1) ※15epochだと精度がよくないので、閾値0.1
run_worflow(im, finetuned_model, oid_labels, 0.1)
# Fine-Tuningしていないモデルで物体検出(閾値0.9)
run_worflow(im, original_model, coco_labels, 0.9)