Safie Engineers' Blog!

Safieのエンジニアが書くブログです

Cleanlabを活用したData-centric AI: ノイズのある物体検出データセットのクレンジング

画像認識

はじめに

セーフィー株式会社 で画像認識AIの開発エンジニアをしている水野です。

現在、AI-App 人数カウントで利用される物体検出モデルの精度改善に取り組んでいます。物体検出モデルの精度改善方法としては様々な手法が考えられますが、近年はData-centric AIというアプローチが注目されています。そこで本稿では、Data-centric AIで物体検出モデルの精度を改善する一手法としてCleanlabを用いたデータセットのクレンジング方法について紹介します。

Data-centric AI

Data-centric AIとは、2021年3月にAndrew Ng氏の講演「From Model-centric to Data-centric AI」で提唱された概念で、AIモデルを固定してデータセットを改善することでモデルの精度改善を実現するアプローチのことです。これは従来主流であったデータセットを固定してAIモデルを改善するModel-centric AIとは対極にある考え方と言えそうです。ただしこれらはどちらのアプローチが優れているという話ではなく、実際の開発では両方の観点で改善していく必要があります。

データセットを改善したい場合、どのようなアプローチが考えられるでしょうか。単純に考えると画像およびアノテーションを人の目で確認し、問題のあるデータを抽出して修正するという方法が挙げられそうですが、この方法にはいくつかの課題があります。

  • 現実のデータセットでは画像枚数が数千枚～数万枚という規模であり、人の目で一通り確認しミス無く修正しきるのはかなり時間が掛かる
  • データセットがクリーンであることを客観的に示すことが難しい

このようなデータセット改善の作業で発生する課題を解決するためにCleanlabのようなData-centric AI向けのライブラリが活用できます。

Cleanlab

Cleanlabはconfident learningの考え方をベースとしてデータセットのノイズを抽出するためのライブラリです。Cleanlabを使用することでデータセット内に存在するアノテーションの問題を自動的に抽出することができます。また画像ごとの品質スコアが出力されるので、アノテーションの品質がどの程度なのか定量的に示すことが可能です。従来は画像識別タスクのみサポートされていたのですが、2023年9月のリリースから物体検出タスクに対応しています。

Cleanlabを用いた基本的なデータクレンジングの流れは下記の通りです。

  1. データセット作成: データセットをtrainとvalidationに分けます
  2. AIモデルの学習: trainデータでモデルを作成します
  3. 予測生成: 学習したモデルを用いてvalidationデータに対する予測を生成します
  4. Cleanlabによる品質スコア算出: validationデータに対する予測と正解データを比較して画像ごとの品質スコアを算出します
  5. アノテーション修正: 品質スコアの低い画像のアノテーションを修正します

お気づきの方もいるかもしれませんが、このままではvalidationデータに含まれる画像の品質しか判定できません。そこで実際にはK-fold分割を用いたOut-of-fold予測を生成し、全データに対する品質スコアを算出します。5-fold分割時の処理を図で書くと以下のようになります。5つに分割したデータで生成した各モデルで予測を生成し、結合することで実質的に全データに対する予測を生成することが出来ます。本稿では簡単化のために技術的な背景や詳細の説明は省略しますが、ご興味のある方は公式の解説ページをご参照ください。

5-fold分割によるOut-of-fold予測生成の流れ (公式の解説ページの図をベースに作成)

物体検出データセットのクレンジング

ここからは物体検出データセットに対する実際のクレンジングの方法を説明します。

物体検出モデルとしてYOLOv8、データセットとしてCOCOを使用します。

データセット作成

クレンジングしたいデータセットをtrainとvalidationに分割します。既に説明した通り、データセット全体をクレンジングする場合はOut-of-fold予測を生成する必要がありますので、COCOデータセット全体をクレンジングしたい場合はCOCOから提供されているtrainデータとvalidationデータを混ぜた上でK-fold分割を実施します。

K-fold分割ではscikit-learnがよく利用されます。データセットの特性に応じて適切な分割方法を選択する必要がありますので、公式ページを参照し各自の状況に合わせた手法を選択してください。また分割数をいくつにすべきかという問題ですが、理想的には多い方が良いのですが分割数が増えると学習に必要な時間が膨大になってしまうため、実用上は5-fold分割がお勧めです。Kaggleのコンペ等でも5-fold分割はよく利用されますし、trainとvalidationの割合としても4:1というのはバランスが良いと思います。

5-fold分割したデータセットの構成例を以下に示します。YOLO形式のラベルフォーマット等は公式の解説ページをご参照ください。またCOCOのフルデータセットの場合、5-fold合計で100GB（1foldあたり20GB）のディスク容量を必要としますのでディスクの空き容量にご注意ください。

$ tree datasets/
datasets/
├── coco_fold1
│   ├── images
│   │   ├── train
│   │   │   └── xxx.jpg
│   │   └── val
│   │       └── yyy.jpg
│   └── labels
│       ├── train
│       │   └── xxx.txt
│       └── val
│           └── yyy.txt
├── coco_fold2
├── coco_fold3
├── coco_fold4
└── coco_fold5

各foldに対応するyamlファイルの作成も必要です。fold1の例を以下に示します。

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco_fold1  # dataset root dir
train: images/train  # train images (relative to 'path')
val: images/val  # val images (relative to 'path')
test:  # test images (optional)

# Classes (80 COCO classes)
names:
  0: person
  1: bicycle
  2: car
  # ...
  77: teddy bear
  78: hair drier
  79: toothbrush

物体検出モデルの学習

K-fold分割で生成した各foldで学習を実行しfold数分のモデルを生成します。まずは学習を実行するために公式のDockerコンテナを立ち上げます。

$ docker run --ipc=host -it --gpus all -v <DATASET_ROOT>:/usr/src/datasets ultralytics/ultralytics:latest

次にさきほど作成した各foldのデータセットごとに学習コマンドを実行します。yoloコマンドで学習を実行する場合のBashでのスクリプト例を示します。YOLOv8ではモデルのバリエーションとしてYOLOv8nからYOLOv8xまで5種類ラインナップがあり、ここではYOLOv8mを使用していますが、Cleanlabはモデルに依存しない手法なので基本的にどのモデルを選択しても問題ありません。実行する計算機環境やどれぐらいの時間でスコアを算出したいかに応じて選択してください。

for i in {1..5}; do
    yolo train data=/usr/src/datasets/coco_fold${i}.yaml model=yolov8m.pt
done

予測と正解データの生成

各foldで学習したモデルで各foldのvalidationデータに対する推論を実行し、各推論結果を結合することでデータセット全体に対する予測を出力します。生成した予測はCleanlabが入力として期待するフォーマットに変換する必要があります。またvalidationデータに対する正解データの生成も必要です。

YOLOv8を用いて予測と正解データの生成を実行するサンプルスクリプトは以下の通りです。出力される predictions.pkl が予測、 labels.pkl が正解データに対応しています。

import os
import pickle

import numpy as np
from tqdm import tqdm
from ultralytics import YOLO
from ultralytics.data.dataset import load_dataset_cache_file
from ultralytics.engine.results import Results

def make_prediction(results: Results, num_classes: int = 80) -> np.ndarray:
    """YOLOv8の推論結果をCleanlabの入力形式に変換"""
    pred_cls = results.boxes.cls.cpu().numpy()
    pred_conf = results.boxes.conf.cpu().numpy()
    pred_boxes = results.boxes.xyxy.cpu().numpy()
    prediction = []
    for target_cls in range(num_classes):
        target_cls_mask = pred_cls == target_cls
        if sum(target_cls_mask) == 0:
            # 該当クラスの検知結果が無い場合の処理
            prediction.append(np.empty((0, 5), dtype=np.float32))
        else:
            target_cls_boxes = pred_boxes[target_cls_mask]
            target_cls_conf = pred_conf[target_cls_mask]
            target_cls_pred = np.concatenate([target_cls_boxes, target_cls_conf[:, None]], axis=1)
            prediction.append(target_cls_pred)
    return np.array(prediction)

def make_label(data: dict) -> dict:
    """YOLOv8の正解データをCleanlabの入力形式に変換"""
    label = {"labels": data["cls"].flatten().astype(int), "seg_map": os.path.basename(data["im_file"])}
    # 正規化されたcx, cy, w, hの2次元アレイをshapeの値で元の座標に戻して、xyxy座標に変換
    bboxes = data["bboxes"] * np.array([data["shape"][1], data["shape"][0], data["shape"][1], data["shape"][0]])
    bboxes[:, 0] = bboxes[:, 0] - bboxes[:, 2] / 2
    bboxes[:, 1] = bboxes[:, 1] - bboxes[:, 3] / 2
    bboxes[:, 2] = bboxes[:, 0] + bboxes[:, 2]
    bboxes[:, 3] = bboxes[:, 1] + bboxes[:, 3]
    label["bboxes"] = bboxes
    return label

def main() -> None:
    """5-foldのOut-of-fold予測および正解データを生成"""
    model_paths = [
        "/usr/src/ultralytics/runs/detect/train/weights/best.pt",
        "/usr/src/ultralytics/runs/detect/train2/weights/best.pt",
        "/usr/src/ultralytics/runs/detect/train3/weights/best.pt",
        "/usr/src/ultralytics/runs/detect/train4/weights/best.pt",
        "/usr/src/ultralytics/runs/detect/train5/weights/best.pt",
    ]
    dataset_paths = [
        "/usr/src/datasets/coco_fold1/labels/val.cache",
        "/usr/src/datasets/coco_fold2/labels/val.cache",
        "/usr/src/datasets/coco_fold3/labels/val.cache",
        "/usr/src/datasets/coco_fold4/labels/val.cache",
        "/usr/src/datasets/coco_fold5/labels/val.cache",
    ]
    predictions: list[np.ndarray] = []
    labels: list[np.ndarray] = []
    for model_path, dataset_path in zip(model_paths, dataset_paths):
        model = YOLO(model_path)
        cache = load_dataset_cache_file(dataset_path)
        for data in tqdm(cache["labels"]):
            results = model(data["im_file"])
            predictions.append(make_prediction(results[0]))
            labels.append(make_label(data))

    # pickleファイルとしてpredictionsとlabelsを保存
    with open("predictions.pkl", "wb") as f:
        pickle.dump(predictions, f)
    with open("labels.pkl", "wb") as f:
        pickle.dump(labels, f)

if __name__ == "__main__":
    main()

Cleanlabによる品質スコア算出

validationデータに対する予測と正解データを比較して画像ごとの品質スコアを算出します。品質スコアは0.0から1.0の範囲の数値で出力され、数値が高いほどアノテーション品質が高いことを表します。先ほどのスクリプトで predictions.pkl と labels.pkl が生成されているはずなので、これをCleanlabの物体検出向けのAPIに入力します。Cleanlabは pip install Cleanlab 等でインストールできます。

スコアを算出し、スコアの低い画像を確認するサンプルコードは以下の通りです。COCOは80クラスのデータセットですが、説明のしやすさのためpersonクラスのみに絞った場合の例になっています。このコードをJupyter Notebook等で実行することで品質スコアの低い画像およびアノテーションを確認することが出来ます。

import pickle
from Cleanlab.object_detection.rank import (
    get_label_quality_scores,
    issues_from_scores,
)
from Cleanlab.object_detection.summary import visualize

IMAGE_PATH = 'all_images'
predictions = pickle.load(open("predictions.pkl", "rb"))
labels = pickle.load(open("labels.pkl", "rb"))

# 各画像の品質スコアを計算
scores = get_label_quality_scores(labels, predictions)

# 品質スコアが0.5を下回る画像インデックスをスコアの低い順にして取得
issue_idx = issues_from_scores(scores, threshold=0.5)

# 一番スコアの低い画像を表示
class_names = {"0": "person"}
issue_to_visualize = issue_idx[0]  # ここの数値を変更することで別の画像を表示可能
label = labels[issue_to_visualize]
prediction = predictions[issue_to_visualize]
image_path = IMAGE_PATH + label['seg_map']    
visualize(image_path, label=label, prediction=prediction, class_names=class_names, overlay=False)

私の手元で試したCOCOデータセットにおけるいくつかスコアの低い画像および高い画像の例を載せておきます。図中の左画像の赤枠が正解データ、右画像の青枠がモデルによる予測です。

アノテーション修正

スコアの計算が出来たらスコアの低い画像群の確認と修正を実施します。アノテーションツールとしてセーフィーではFastLabel社のツールをよく利用しています。その他CVATやYOLOv8公式でサポートされているRoboflow等、様々なツールがありますので利用しやすいものを使用してください。

Cleanlabをフル活用するためのTips

今回Cleanlabを利用する中で気づいた課題やそれに対する対策についていくつか説明します。

同じ正解データや予測を複数回マッチングさせないようにする

Cleanlabの現状の実装（2024/05/17時点）ではスコア計算時に同じ正解データや予測が複数回マッチングされる場合があり、スコアが正しく計算できない場合があります。スコアを出力してみて複数回マッチングの影響で意図しないスコアが出ているようであれば、スコア計算処理のマッチング部分を修正する必要があります。

閾値をモデルやデータセットに応じて適切に設定する

Cleanlabでは信頼度の低い予測をスコア計算に使用しないように2つの閾値を設けています。この閾値の設定によって出力されるスコアの傾向が大きく変わるので、自分が学習したモデルやデータセットの特性に応じて調整することで精度の高いスコア計算が可能になります。

スコアが1.0になっている画像でも一通り確認する

Cleanlabのスコアは0.0から1.0までの数値で表されるため、定義としては1.0は最もスコアが高い（アノテーション品質が高い）ことになります。しかし正解データと予測でマッチングするのものが無かった場合等にCleanlabはスコアのデフォルト値である1.0を出力するため、1.0というスコアは必ずしもアノテーションに誤りが無いことを表してはいません。実際上正解データと予測がぴったり合う（スコアが1.0になる）ことはほとんどあり得ないので、1.0のスコアが出力されたということは何らかの問題がある可能性があります。

以下の画像は一例ですが、スコアが1.0にも関わらず正解データではアノテーションが1つ漏れていることが分かります。これは予測のbboxの信頼度が低く、正解データとのマッチングプロセスがスキップされるために発生していると考えられます。

この画像は Brad Greenlee氏によるDonut Tower II をCleanlabにより可視化したものです。
(C) 2007 Brad Greenlee, Donut Tower II, License:http://creativecommons.org/licenses/by/2.0/

まとめ

今回はCleanlabを用いたData-centric AIによる物体検出データセットのクレンジング方法について紹介しました。Cleanlabを利用することで効率的にデータセットを改善することが出来、モデルを変えなくてもモデルの精度改善が可能になります。実際の改善では今回紹介したData-centricおよびModel-centricの両方の観点で取り組むことが重要ではありますが、Model-centricアプローチによる精度改善に手詰まりを感じられている方がいれば試してみるのも良いと思います。

最後になりますが、セーフィーではエンジニアを積極的に募集しています。気になる方はこちらをご覧ください！

https://safie.co.jp/teams/engineering/

カジュアル面談から受け付けておりますので、気軽に応募いただければと思います！

最後までお読みいただき、ありがとうございました。

© Safie Inc.