MMDetectionでCoDETRをトレーニングする

# MMDetectionでモデルを学習させる方法

たくさんモデルが使えて、State of art（現状最高）のモデルとかもMMDetectionで動いているで有名なMMDetectionですが、
バージョンアップですぐ動かなくなることでも僕の中で有名です。

**この記事では、ちゃんとMMDetectionを動かして、自分のデータでモデルを学習させます。**

# Install
### たのむよ
まあバージョンアップごとにいろんなエラーが出て、
うまくうごかないんだけど、
とりあえずこれでうごいたで。
2024/11/16の時点で。

### 動くセットアップ

インストール。
**mimではなくpipでインストールすることと、mmcvのバージョンがポイント。**

```shell
pip install mmcv==2.1.0
git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
pip install -v -e .
```

インストールに３０分以上かかるけどめげずにがんばろう。

mmdetection/hack_registry.pyというファイルを作る。

```hack_registry.py
#  hack_registry.py
import logging

from mmengine.registry import Registry
from mmengine.logging import print_log
from typing import Type, Optional, Union, List


def _register_module(self,
                     module: Type,
                     module_name: Optional[Union[str, List[str]]] = None,
                     force: bool = False) -> None:
    """Register a module.

    Args:
        module (type): Module to be registered. Typically a class or a
            function, but generally all ``Callable`` are acceptable.
        module_name (str or list of str, optional): The module name to be
            registered. If not specified, the class name will be used.
            Defaults to None.
        force (bool): Whether to override an existing class with the same
            name. Defaults to False.
    """
    if not callable(module):
        raise TypeError(f'module must be Callable, but got {type(module)}')

    if module_name is None:
        module_name = module.__name__
    if isinstance(module_name, str):
        module_name = [module_name]
    for name in module_name:
        if not force and name in self._module_dict:
            existed_module = self.module_dict[name]
            # raise KeyError(f'{name} is already registered in {self.name} '
            #                f'at {existed_module.__module__}')
            print_log(
                f'{name} is already registered in {self.name} '
                f'at {existed_module.__module__}. Registration ignored.',
                logger='current',
                level=logging.INFO
            )
        self._module_dict[name] = module


Registry._register_module = _register_module
```

mmdetのimportの前にhack_registryをimport。

```python
import hack_registry

from mmdet.apis import DetInferencer
```

これでやっと動く。
たのむよ、mmlabさん。

# データセットの用意
COCO形式のデータセットを用意する。
私はYolo形式のデータセットからCOCO形式に変換したので、以下のスクリプトを用いました。
他の形式の方はがんばってください😉

データの移動には以下のコマンドが高速です。shutilだと大量のデータは重い。

```python
source_dir = "/content/drive/MyDrive/datasets/defect_crop_car/labels/train"
destination_dir = "/content/defect_crop_car/labels/train"  

!mkdir -p "{destination_dir}"

!find "{source_dir}" -type f | xargs -P 8 -I{{}} cp {{}} "{destination_dir}"
```

```python
import json
import collections as cl
import cv2

def get_info():
  tmp = cl.OrderedDict()
  tmp["description"] = "my dataset"
  tmp["url"] = ""
  tmp["version"] = "1"
  tmp["year"] = 2024
  tmp["contributor"] = ""
  tmp["data_created"] = "2024/11/17"
  return tmp

def get_licenses():
  tmp = cl.OrderedDict()
  tmp["id"] = 0
  tmp["url"] = ""
  tmp["name"] = ""
  return [tmp]

def get_image_data(id, image_path, image_file, w, h):
  tmp = cl.OrderedDict()
  tmp["license"] = 0
  tmp["id"] = id
  tmp["file_name"] = image_file
  tmp["width"] = w
  tmp["height"] = h
  tmp["date_captured"] = ""
  tmp["coco_url"] = ""
  tmp["flickr_url"] = ""
  return tmp

def get_annotation(id, image_id, category_id, bbox):
  tmp = cl.OrderedDict()

  tmp_segmentation = cl.OrderedDict()
  tmp_segmentation = [bbox]
  tmp["segmentation"] = tmp_segmentation
  tmp["id"] = id
  tmp["image_id"] = image_id
  tmp["category_id"] = category_id
  tmp["area"] = bbox[2]*bbox[3]
  tmp["iscrowd"] = 0
  tmp["bbox"] =  bbox
  return tmp

def get_categories():
  tmps = []
  sup = ["damage"]
  classes = ["your_class0", "your_class1"]
  for i in range(len(classes)):
    tmp = cl.OrderedDict()
    tmp["id"] = i
    tmp["supercategory"] = sup
    tmp["name"] = classes[i]
    tmps.append(tmp)
  return tmps

def create_dataset(image_dir, label_dir, json_path):
    image_list = os.listdir(image_dir)
    print(len(image_list))
    info = get_info()
    licenses = get_licenses()
    categories = get_categories()
    images = []
    annotations = []
    
    for i, image_file in enumerate(image_list):
    # for i in range(100):
      image_file = image_list[i]
      image_path = os.path.join(image_dir,image_file)
      label_file = os.path.splitext(os.path.basename(image_file))[0] + ".txt"
      label_path = os.path.join(label_dir,label_file)
      if os.path.exists(label_path):
        img = cv2.imread(image_path)
        img_h, img_w, _ = img.shape
    
        image_data = get_image_data(i, image_path, image_file, img_w, img_h)
        source_file = open(label_path)
        for defect_index, line in enumerate(source_file):
          staff = line.split()
          class_idx = int(staff[0])
    
          x_center, y_center, width, height = float(staff[1])*img_w, float(staff[2])*img_h, float(staff[3])*img_w, float(staff[4])*img_h
          x = round(x_center-width/2,2)
          y = round(y_center-height/2,2)
          width = round(width,2)
          height = round(height,2)
          bbox = [x, y, width, height]
          id = i * 1000 + defect_index
          annotation = get_annotation(id, i, class_idx, bbox)
          annotations.append(annotation)
          images.append(image_data)
    
    json_data = {
        'info': info,
        'images': images,
        'licenses': licenses,
        'annotations': annotations,
        'categories': categories,
    }
    
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(json_data, f, ensure_ascii=False)
```

```python
import os

dataset_root_dir = "my_dataset/"
image_dir = os.path.join(dataset_root_dir,"images/")
label_dir = os.path.join(dataset_root_dir,"labels/")

image_train_dir = os.path.join(image_dir, "train")
image_val_dir = os.path.join(image_dir, "val")
label_train_dir = os.path.join(label_dir, "train")
label_val_dir = os.path.join(label_dir, "val")

dataset_name = os.path.basename(os.path.normpath(dataset_root_dir))
train_json_path = os.path.join(dataset_root_dir, "train.json")
val_json_path = os.path.join(dataset_root_dir, "val.json")

create_dataset(image_train_dir, label_train_dir, train_json_path)
create_dataset(image_val_dir, label_val_dir, val_json_path)
```

これでdatasetのディレクトリにtrain.json, val.jsonができます。

### configの作成と編集

```python:my_config.py
_base_ = '/content/mmdetection/projects/CO-DETR/configs/codino/co_dino_5scale_swin_l_16xb1_1x_coco.py' 

classes = ("class0","class1")
data_root = 'my_dataset'

train_dataloader = dict(
    batch_size=1,
    dataset=dict(
        ann_file='train.json',
        data_prefix=dict(img='/images/train/')
    ),
)
val_dataloader = dict(
    batch_size=1,
    dataset=dict(
        ann_file='my_dataset/val.json',
        data_prefix=dict(img='images/val/')
    ),
)

test_dataloader = dict(
    batch_size=1,
    dataset=dict(
        ann_file='my_dataset/val.json',
        data_prefix=dict(img='images//val/')
    ),
)

test_evaluator = dict(
    ann_file='my_dataset/val.json',
)

val_evaluator = dict(
    ann_file='my_dataset/val.json',
)

load_from = 'co_dino_5scale_swin_large_1x_coco-27c13da4.pth'
```

mmdetection/mmdet/dataset/coco.py
のCocoDatasetのMETAINFOをカスタムデータセットのクラスに置き換える。

```python:coco.py
@DATASETS.register_module()
class CocoDataset(BaseDetDataset):
    """Dataset for COCO."""

    METAINFO = {
      'classes': ("class0", "class1"),
      'palette': [(220, 20, 60), (119, 11, 32)]
    }
```

使用するconfigの_base_ configを辿って、num_classesを書き換える。
僕の例では、

_base_ = '/content/mmdetection/projects/CO-DETR/configs/codino/co_dino_5scale_swin_l_16xb1_1x_coco.py' 
のconfig ファイルのさらに_base_が
/content/mmdetection/projects/CO-DETR/configs/codino/co_dino_5scale_r50_lsj_8xb2_1x_coco.py
なので、それのnum_classesを自分のデータセットのクラス数にします。

```python:co_dino_5scale_r50_lsj_8xb2_1x_coco.py
num_classes = 2
```

これで学習を開始できます。

```shell
python tools/train.py my_config.py --work-dir /content/drive/MyDrive/codino
```

--work-dirに指定したパスに自動で結果保存dirが作成されます。
この場合、codinoは実行スクリプトでmakedirされます。

データ数が多い時は、小量のテストデータを作ってトレーニングが正しく１エポック実行されるか確認してから全データを学習すると良いと思います。全データで1epoch待ってエラー出たら解決に時間がかかる。




🐣
****
フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

機械学習関連の情報を発信しています。

[Twitter](https://twitter.com/JackdeS11)
[Medium](https://medium.com/@rockyshikoku)
[GitHub](https://github.com/john-rocky)