Training YOLO with custom dataset by fiftyone

Labelme로 만들어낸 커스텀 데이터셋을 이용해서 이제 yolo를 트레이닝 시킬 차례다.

이 과정에 fiftyone 이라는 툴을 사용해볼 예정

자 그럼 시작해보자

지난 포스팅에서 fiftyone 상에 COCO 포맷으로 뽑아둔 데이터셋이 잘 올라가는 것까지 확인을 했다.

이제 이걸 yolo 학습용 데이터셋으로 변경해야한다.

일단 이미지 수가 매우 적긴 하지만, 작동이 되는지 확인을 하기 위한 용도로 진행해보면,

import fiftyone as fo
import fiftyone.utils.random as four

def export_yolo_data(
    samples,
    export_dir,
    classes,
    label_field = "detections",
    split = None
    ):

    if type(split) == list:
        splits = split
        for split in splits:
            export_yolo_data(
                samples,
                export_dir,
                classes,
                label_field,
                split
            )
    else:
        if split is None:
            split_view = samples
            split = "val"
        else:
            split_view = samples.match_tags(split)

        split_view.export(
            export_dir=export_dir,
            dataset_type=fo.types.YOLOv5Dataset,
            label_field=label_field,
            classes=classes,
            split=split
        )

# The directory containing the source images
data_path = "data_dataset_coco_weldHelmet"

# The path to the COCO labels JSON file
labels_path = "data_dataset_coco_weldHelmet/annotations.json"

# Import the dataset
dataset = fo.Dataset.from_dir(
    dataset_type=fo.types.COCODetectionDataset,
    data_path=data_path,
    labels_path=labels_path,
)

dataset.untag_samples(dataset.distinct("tags"))

classes=["WeldHelmet"]

## split into train and val
four.random_split(
    dataset,
    {"train": 0.8, "val": 0.2}
)

## export in YOLO format
export_yolo_data(
    dataset,
    "WeldHelmet_train",
    classes,
    split = ["train", "val"]
)

8번쨰 줄을 ground_truth에서 detections로 수정했다.

이걸 하고 나면