🐈
Bounding Boxの自動ラベリングの試み(Detectron2)
やりたいこと
- 精度強めのモデルである程度の量でFTしたら、自動ラベリングできないか。
- Bboxのアノテーションは修正も楽だから、一度ラベル済のをツールでパラパラめくって、間違っているものを手動ラベリングすることで効率上がらないか。
- Detecton2を触ってみたい、あ~ツールキットね~って言えるようになりたい!
※Twitterやコメントなどで諸々は優しくアドバイスやご指摘などください。
Detectron2って
Detectron2はFacebook AIのオープンソースPJで、PyTorchの物体検出ツールキットです。
検出できる内容はざっくり以下。
- 物体検知 Bounding Box(Faster R-CNN/RetinaNet)
- インスタンスセグメンテーション(Mask R-CNN)
- 姿勢推定(Keypoint_RCNN)
- パノプティックセグメンテーション(Panoptic FPN)
MMDetectionとかBigDetectionとか他にも色々あるようですが、
お師匠さん使ってるようなので。(そこまで最新ではない?)
数モデルしかないからYOLOとか使いたい人とかは他の使うんじゃないかなと思います。
環境構築
git clone https://github.com/facebookresearch/detectron2.git
pip install -e detectron2
でいける人はいける。ですが自分はWindowsのローカルでtorchとかライブラリ周りで若干ガチャガチャした記憶が。dockerとか考えればもっと良い記事が書けた?
学習
今回はスターウォーズのドロイド達(C3PO/R2D2/BB-8,DTKCは含まず)の写真を50枚学習。
アノテはcoco-annotatorを今回使ってみました。矩形の修正する動作がやり方分からなかったですが、管理画面もあって、docker-composeで速攻立ち上がるので使いやすかったです。
model_zooのDetectionから、faster_rcnn_X_101_32x8d_FPN_3が一番精度は高いんじゃないかなーと選びました。
train.py
import os
import json
import cv2
import random
import torch
import detectron2
import numpy as np
from detectron2 import model_zoo
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.data.datasets import load_coco_json
from detectron2.engine import DefaultPredictor, DefaultTrainer
from detectron2.config import get_cfg
from detectron2.utils.logger import setup_logger
from detectron2.utils.visualizer import Visualizer
CLASSES_TXT_PATH = "./img_train/classes.txt"
IMG_TRAIN_JSON_PATH = "./img_train/img_train.json"
IMG_TRAIN_PATH = "./img_train"
TRAIN_DATASET_NAME = "img_train"
MODEL_CONFIG_PATH = "COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"
CKPT_FOLDER_PATH = os.path.join(os.getcwd(), "ckpt/img_train")
def get_class_names(txt_path):
with open(txt_path, 'r') as f:
class_names = [line.strip() for line in f.readlines()]
return class_names
class_names = get_class_names(CLASSES_TXT_PATH)
coins_metadata = MetadataCatalog.get(TRAIN_DATASET_NAME).set(thing_classes=class_names) # Set metadata here
dataset_dicts = load_coco_json(IMG_TRAIN_JSON_PATH, IMG_TRAIN_PATH, TRAIN_DATASET_NAME)
DatasetCatalog.register(TRAIN_DATASET_NAME, lambda: dataset_dicts) # Register dataset
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(MODEL_CONFIG_PATH))
cfg.DATASETS.TRAIN = (TRAIN_DATASET_NAME,)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 0
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(MODEL_CONFIG_PATH)
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.0004
cfg.SOLVER.MAX_ITER = 1500
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(class_names)
os.makedirs(CKPT_FOLDER_PATH, exist_ok=True)
cfg.OUTPUT_DIR = CKPT_FOLDER_PATH
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
setup_logger()
推論
では学習させた重みを利用して、推論させてlabelsフォルダにラベルも吐かせようと思います。pascal_vocのラベルは正直あってるか分かりません。。。今回使ったテスト用画像と学習用画像はzipファイルでアップしてます(良いのか？)。※マイリポジトリ
inference.py
import os
import glob
import cv2
from detectron2 import model_zoo
from detectron2.utils.logger import setup_logger
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.utils.visualizer import Visualizer
from label_saver import get_class_names, get_image_files, save_coco_annotation, save_yolo_label, save_pascal_voc_annotation, convert_to_yolo_label
setup_logger()
MODEL_WEIGHTS_PATH = os.path.join("ckpt/img_train", "model_final.pth")
CLASS_NAMES_PATH = "./img_train/classes.txt"
IMAGE_FOLDER = "./img_test"
COCO_ANNOTATIONS_OUTPUT_FOLDER = "labels"
CONFIDENCE_THRESHOLD = 0.7
def detect_and_visualize(image_path, predictor, cfg, class_names, save_yolo_label, save_pascal_voc_annotation):
im = cv2.imread(image_path)
if im is None:
print(f"画像ファイルの読み込みに失敗しました: {image_path}")
return
# coco dataset
# v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
# custom dataset
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get("img_train"), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
outputs = predictor(im)
# 画像を確認する
# cv2.imshow("Object Detection", v.get_image()[:, :, ::-1])
# cv2.waitKey(0)
instances = outputs['instances']
height, width = im.shape[:2]
annotations = []
for idx, bbox in enumerate(instances.pred_boxes):
class_id = instances.pred_classes[idx].item()
x1, y1, x2, y2 = bbox.tolist()
width, height = x2 - x1, y2 - y1
coco_bbox = [x1, y1, width, height]
annotations.append({"class_id": class_id, "bbox": coco_bbox})
# YOLOアノテーションの保存
x_center, y_center, w, h = convert_to_yolo_label(bbox, width, height)
save_yolo_label(image_path, class_id, x_center, y_center, w, h)
# Save Pascal VOC annotations
save_pascal_voc_annotation(image_path, class_id, x1, y1, x2, y2, class_names)
return {"image_path": image_path, "annotations": annotations}
def main():
cfg = get_cfg()
# pre-trained model
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")]
# custom dataset
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = CONFIDENCE_THRESHOLD
cfg.MODEL.WEIGHTS = MODEL_WEIGHTS_PATH
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
class_names = get_class_names(CLASS_NAMES_PATH)
MetadataCatalog.get("img_train").set(thing_classes=class_names)
predictor = DefaultPredictor(cfg)
all_annotations = []
image_files = get_image_files(IMAGE_FOLDER)
for image_file in image_files:
annotations = detect_and_visualize(image_file, predictor, cfg, class_names, save_yolo_label, save_pascal_voc_annotation)
if annotations is not None:
all_annotations.append(annotations)
save_coco_annotation(IMAGE_FOLDER, all_annotations, COCO_ANNOTATIONS_OUTPUT_FOLDER, class_names)
cv2.destroyAllWindows()
if __name__ == "__main__":
main()
で細かいラベル保存などの関数はこちらのファイルから読み込ませます。
label_saver.py
import os
import glob
import json
import cv2
import random
import xml.etree.ElementTree as ET
from pathlib import Path
from detectron2.utils.visualizer import Visualizer
from detectron2.data import DatasetCatalog, MetadataCatalog
YOLO_LABELS_OUTPUT_FOLDER = "labels/yolo"
PASCAL_VOC_ANNOTATIONS_OUTPUT_FOLDER = "labels/pascal_voc"
def get_class_names(txt_path):
with open(txt_path, 'r') as f:
class_names = [line.strip() for line in f.readlines()]
return class_names
def get_image_files(folder_path):
extensions = ['jpg', 'jpeg', 'png']
return [f for ext in extensions for f in glob.glob(os.path.join(folder_path, f'*.{ext}'))]
def save_coco_annotation(image_folder, all_annotations, output_folder, class_names):
coco_annotations = {
"images": [],
"annotations": [],
"categories": [
{"id": idx, "name": name} for idx, name in enumerate(class_names)
],
}
annotation_id = 0
for item in all_annotations:
image_path = item["image_path"]
image_id = Path(image_path).stem
image = cv2.imread(image_path)
height, width = image.shape[:2]
coco_annotations["images"].append({
"id": image_id,
"dataset_id": 1, # or other dataset id you want to use
"category_ids": [annotation["class_id"] for annotation in item["annotations"]],
"path": image_path,
"width": width,
"height": height,
"file_name": os.path.basename(image_path),
"annotated": True,
"annotating": [],
"num_annotations": len(item["annotations"]),
"metadata": {},
"milliseconds": 0,
"events": [],
"regenerate_thumbnail": False,
"is_modified": False
})
for annotation in item["annotations"]:
coco_annotation = {
"id": annotation_id,
"image_id": image_id,
"category_id": annotation["class_id"],
"bbox": annotation["bbox"],
"segmentation": [],
"iscrowd": 0,
}
coco_annotations["annotations"].append(coco_annotation)
annotation_id += 1
output_path = os.path.join(output_folder, "coco.json")
with open(output_path, "w") as f:
json.dump(coco_annotations, f, indent=2)
def save_yolo_label(image_path, class_id, x_center, y_center, w, h):
label_path = os.path.join(os.path.dirname(os.path.abspath(image_path)), "..", YOLO_LABELS_OUTPUT_FOLDER)
os.makedirs(label_path, exist_ok=True)
label_file = os.path.join(label_path, os.path.splitext(os.path.basename(image_path))[0] + ".txt")
with open(label_file, "a") as f:
f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n")
def convert_to_yolo_label(box, width, height):
x_center = (box[0] + box[2]) / 2 / width
y_center = (box[1] + box[3]) / 2 / height
w = (box[2] - box[0]) / width
h = (box[3] - box[1]) / height
return x_center, y_center, w, h
def save_pascal_voc_annotation(image_path, class_id, x1, y1, x2, y2, class_names):
image_folder_path = os.path.dirname(os.path.abspath(image_path))
annotation_folder = os.path.join(image_folder_path, "..", PASCAL_VOC_ANNOTATIONS_OUTPUT_FOLDER)
os.makedirs(annotation_folder, exist_ok=True)
annotation_file = os.path.join(annotation_folder, os.path.splitext(os.path.basename(image_path))[0] + ".xml")
if not os.path.exists(annotation_file):
create_pascal_voc_xml(image_path, annotation_file)
tree = ET.parse(annotation_file)
root = tree.getroot()
obj = ET.SubElement(root, "object")
ET.SubElement(obj, "name").text = class_names[class_id]
ET.SubElement(obj, "pose").text = "Unspecified"
ET.SubElement(obj, "truncated").text = "0"
ET.SubElement(obj, "difficult").text = "0"
bndbox = ET.SubElement(obj, "bndbox")
ET.SubElement(bndbox, "xmin").text = str(x1)
ET.SubElement(bndbox, "ymin").text = str(y1)
ET.SubElement(bndbox, "xmax").text = str(x2)
ET.SubElement(bndbox, "ymax").text = str(y2)
tree.write(annotation_file)
def create_pascal_voc_xml(image_path, annotation_file):
img = cv2.imread(image_path)
height, width, depth = img.shape
root = ET.Element("annotation")
folder = ET.SubElement(root, "folder")
folder.text = "images"
filename = ET.SubElement(root, "filename")
filename.text = os.path.basename(image_path)
source = ET.SubElement(root, "source")
ET.SubElement(source, "database").text = "Unknown"
size = ET.SubElement(root, "size")
ET.SubElement(size, "width").text = str(width)
ET.SubElement(size, "height").text = str(height)
ET.SubElement(size, "depth").text = str(depth)
segmented = ET.SubElement(root, "segmented")
segmented.text = "0"
tree = ET.ElementTree(root)
tree.write(annotation_file)
結果
20枚程簡易テストしましたが、SW公式ドロイドは全て検出、人間の誤検出は無し、まあまあ行けてました。また、他のロボットもほぼドロイド判定をもらう中、ガンダムには反応しませんでした。もしかしたら深い哲学を獲得したのかもしれません。
-
上手くできている例
-
誤検出例(後ろのボトル?やむなし)
感想 + その他
- GPTラベリングが外注を超える精度が～とか、diffusionにデータセット作らせるみたいな話も聞いたりする中、CV勢の皆様が普段どうされているのか気になります。製品化実務だとちゃんとデータセット時間かけてこうぜになりやすいみたいなツイートも見てそりゃそうですよねーと思いましたが、半教師学習とかもzennにはそんな記事無さそうだったので書いてみたいですね～。
- zennはとっても書きやすい、wpの自ブログは廃止じゃ廃止。ソースコードはもっと引数使う感じにかっこよく書けるようになりたいですね。
- tensorboardでログも一応見れます。まあ今回は見てもそんなに参考になりませんが。
-
マイリポジトリ
apiでCC0ダウンロードしてくるスクリプトも一応中にありますが、キーワードで欲しいものが
集まるかはよりけりです。あとpexelsにあるR2D2とは一体。。。
Discussion