在detectron2 上注册 yolo 格式数据集

注册coco格式数据集

detectron上的coco数据集的注册方式如下：

from detectron2.data.datasets import register_coco_instances
register_coco_instances("my_dataset_train", {}, 'PathofTrainJson.json', 'PathofTrainImage')
register_coco_instances("my_dataset_test", {}, 'PathofTestJson.json', 'PathofTestImage')

注册yolo格式

注意不是yolo-obb格式，yolo格式的标注如下：

yolo转coco的代码如下，其中get_dicts是把yolo格式转换coco格式，注册的关键是DatasetCatalog.register：

from detectron2.engine import DefaultTrainer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.structures import BoxMode
from detectron2.config import get_cfg as _get_cfg
from detectron2 import model_zoo

def register_datasets(root_dir, class_list_file):
    """
    Registers the train and validation datasets and returns the number of classes.

    Args:
        root_dir (str): Path to the root directory of the dataset.
        class_list_file (str): Path to the file containing the list of class names.

    Returns:
        int: The number of classes in the dataset.
    """
    # Read the list of class names from the class list file.
    with open(class_list_file, 'r') as reader:
        classes_ = [l[:-1] for l in reader.readlines()]

    # Register the train and validation datasets.
    for d in ['train', 'valid']:  #'train' 和 'valid'分别为数据集注册的名字，可更改
        DatasetCatalog.register(d, lambda d=d: get_dicts(os.path.join(root_dir, d, 'images'),
                                                         os.path.join(root_dir, d, 'labels')))
        # Set the metadata for the dataset.
        MetadataCatalog.get(d).set(thing_classes=classes_)

    return len(classes_)

def get_dicts(img_dir, ann_dir):
    """
    Read the annotations for the dataset in YOLO format and create a list of dictionaries containing information for each
    image.

    Args:
        img_dir (str): Directory containing the images.
        ann_dir (str): Directory containing the annotations.

    Returns:
        list[dict]: A list of dictionaries containing information for each image. Each dictionary has the following keys:
            - file_name: The path to the image file.
            - image_id: The unique identifier for the image.
            - height: The height of the image in pixels.
            - width: The width of the image in pixels.
            - annotations: A list of dictionaries, one for each object in the image, containing the following keys:
                - bbox: A list of four integers [x0, y0, w, h] representing the bounding box of the object in the image,
                        where (x0, y0) is the top-left corner and (w, h) are the width and height of the bounding box,
                        respectively.
                - bbox_mode: A constant from the `BoxMode` class indicating the format of the bounding box coordinates
                             (e.g., `BoxMode.XYWH_ABS` for absolute coordinates in the format [x0, y0, w, h]).
                - category_id: The integer ID of the object's class.
    """
    dataset_dicts = []
    for idx, file in enumerate(os.listdir(ann_dir)):
        # annotations should be provided in yolo format

        record = {}

        # filename = os.path.join(img_dir, file[:-4] + '.jpg')
        filename = os.path.join(img_dir, file[:-4] + '.jpg')
        height, width = cv2.imread(filename).shape[:2]
        # print('$'+str(idx))
        # print('#'+str(file))
        record["file_name"] = filename
        record["image_id"] = idx
        record["height"] = height
        record["width"] = width

        objs = []
        with open(os.path.join(ann_dir, file)) as r:
             #修复一行是 末尾被删除的情况
            # lines = [l[:-1] for l in r.readlines()]
            lines = [l for l in r.readlines()]
            lens = len(lines)
            if lens > 1 :
                for i in range(0,lens-1):
                    lines[i] = lines[i][:-1]

                # print("@"+str(lines))

        for _, line in enumerate(lines):
            if len(line) > 2:
                parts = line.split(' ')   
                if len(parts) == 5:
                    label, cx, cy, w_, h_ = parts
                else:
                    print("Expected 5 parts, got {} at {}".format(len(parts),file))
                # print(file+'  '+'$'+label+' '+cx+' '+cy+' '+w_+' '+h_)
                obj = {
                    "bbox": [int((float(cx) - (float(w_) / 2)) * width),
                             int((float(cy) - (float(h_) / 2)) * height),
                             int(float(w_) * width),
                             int(float(h_) * height)],
                    "bbox_mode": BoxMode.XYWH_ABS,
                    "category_id": int(label),
                }

                objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
        # print(str(dataset_dicts) + '\n')
    return dataset_dicts

之后在使用时调用register_datasets(data_dir, class_list_file) 就可注册，此处的class_list_file指向存放类名的文档（class.name）路径,class.name内容示例，行标应与该类的标注id对应：

调用注册的数据验证目标检测模型

以下附上验证模型的示例代码：

from util import register_datasets
from detectron2.data import build_detection_test_loader
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor

# from detectron2.modeling import build_model
# from detectron2.checkpoint import DetectionCheckpointer

class_list_file = '/pathof/class.names'
data_dir = '/pathof/myDataSet3'
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file('COCO-Detection/retinanet_R_101_FPN_3x.yaml'))
cfg.MODEL.WEIGHTS = './output2/model_0002999.pth' #你训练出的某型

predictor = DefaultPredictor(cfg)
register_datasets(data_dir, class_list_file)
evaluator = COCOEvaluator("valid", output_dir="./output2/")
val_loader = build_detection_test_loader(cfg, "valid") #调用注册后的验证数据集
inference_on_dataset(predictor.model, val_loader, evaluator)

 

参考：

Detectron2 Tutorial.ipynb - Colab (google.com)

Use Custom Datasets — detectron2 0.6 documentation

computervisioneng/train-object-detector-detectron2 (github.com)