third_party/mmyolo/tools/dataset_converters/yolo2coco.py

"""This script helps to convert yolo-style dataset to the coco format.

Usage:
    $ python yolo2coco.py /path/to/dataset # image_dir

Note:
    1. Before running this script, please make sure the root directory
    of your dataset is formatted in the following struction:
    .
    └── $ROOT_PATH
        ├── classes.txt
        ├── labels
        │    ├── a.txt
        │    ├── b.txt
        │    └── ...
        ├── images
        │    ├── a.jpg
        │    ├── b.png
        │    └── ...
        └── ...
    2. The script will automatically check whether the corresponding
    `train.txt`, ` val.txt`, and `test.txt` exist under your `image_dir`
    or not. If these files are detected, the script will organize the
    dataset. The image paths in these files must be ABSOLUTE paths.
    3. Once the script finishes, the result files will be saved in the
    directory named 'annotations' in the root directory of your dataset.
    The default output file is result.json. The root directory folder may
    look like this in the root directory after the converting:
    .
    └── $ROOT_PATH
        ├── annotations
        │    ├── result.json
        │    └── ...
        ├── classes.txt
        ├── labels
        │    ├── a.txt
        │    ├── b.txt
        │    └── ...
        ├── images
        │    ├── a.jpg
        │    ├── b.png
        │    └── ...
        └── ...
    4. After converting to coco, you can use the
    `tools/analysis_tools/browse_coco_json.py` script to visualize
    whether it is correct.
"""
import argparse
import os
import os.path as osp

import mmcv
import mmengine

IMG_EXTENSIONS = ('.jpg', '.png', '.jpeg')


def check_existence(file_path: str):
    """Check if target file is existed."""
    if not osp.exists(file_path):
        raise FileNotFoundError(f'{file_path} does not exist!')


def get_image_info(yolo_image_dir, idx, file_name):
    """Retrieve image information."""
    img_path = osp.join(yolo_image_dir, file_name)
    check_existence(img_path)

    img = mmcv.imread(img_path)
    height, width = img.shape[:2]
    img_info_dict = {
        'file_name': file_name,
        'id': idx,
        'width': width,
        'height': height
    }
    return img_info_dict, height, width


def convert_bbox_info(label, idx, obj_count, image_height, image_width):
    """Convert yolo-style bbox info to the coco format."""
    label = label.strip().split()
    x = float(label[1])
    y = float(label[2])
    w = float(label[3])
    h = float(label[4])

    # convert x,y,w,h to x1,y1,x2,y2
    x1 = (x - w / 2) * image_width
    y1 = (y - h / 2) * image_height
    x2 = (x + w / 2) * image_width
    y2 = (y + h / 2) * image_height

    cls_id = int(label[0])
    width = max(0., x2 - x1)
    height = max(0., y2 - y1)
    coco_format_info = {
        'image_id': idx,
        'id': obj_count,
        'category_id': cls_id,
        'bbox': [x1, y1, width, height],
        'area': width * height,
        'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]],
        'iscrowd': 0
    }
    obj_count += 1
    return coco_format_info, obj_count


def organize_by_existing_files(image_dir: str, existed_categories: list):
    """Format annotations by existing train/val/test files."""
    categories = ['train', 'val', 'test']
    image_list = []

    for cat in categories:
        if cat in existed_categories:
            txt_file = osp.join(image_dir, f'{cat}.txt')
            print(f'Start to read {cat} dataset definition')
            assert osp.exists(txt_file)

            with open(txt_file) as f:
                img_paths = f.readlines()
                img_paths = [
                    os.path.split(img_path.strip())[1]
                    for img_path in img_paths
                ]  # split the absolute path
                image_list.append(img_paths)
        else:
            image_list.append([])
    return image_list[0], image_list[1], image_list[2]


def convert_yolo_to_coco(image_dir: str):
    """Convert annotations from yolo style to coco style.

    Args:
        image_dir (str): the root directory of your datasets which contains
            labels, images, classes.txt, etc
    """
    print(f'Start to load existing images and annotations from {image_dir}')
    check_existence(image_dir)

    # check local environment
    yolo_label_dir = osp.join(image_dir, 'labels')
    yolo_image_dir = osp.join(image_dir, 'images')
    yolo_class_txt = osp.join(image_dir, 'classes.txt')
    check_existence(yolo_label_dir)
    check_existence(yolo_image_dir)
    check_existence(yolo_class_txt)
    print(f'All necessary files are located at {image_dir}')

    train_txt_path = osp.join(image_dir, 'train.txt')
    val_txt_path = osp.join(image_dir, 'val.txt')
    test_txt_path = osp.join(image_dir, 'test.txt')
    existed_categories = []
    print(f'Checking if train.txt, val.txt, and test.txt are in {image_dir}')
    if osp.exists(train_txt_path):
        print('Found train.txt')
        existed_categories.append('train')
    if osp.exists(val_txt_path):
        print('Found val.txt')
        existed_categories.append('val')
    if osp.exists(test_txt_path):
        print('Found test.txt')
        existed_categories.append('test')

    # prepare the output folders
    output_folder = osp.join(image_dir, 'annotations')
    if not osp.exists(output_folder):
        os.makedirs(output_folder)
        check_existence(output_folder)

    # start the convert procedure
    with open(yolo_class_txt) as f:
        classes = f.read().strip().split()

    indices = os.listdir(yolo_image_dir)
    total = len(indices)

    dataset = {'images': [], 'annotations': [], 'categories': []}
    if existed_categories == []:
        print('These files are not located, no need to organize separately.')
        for i, cls in enumerate(classes, 0):
            dataset['categories'].append({'id': i, 'name': cls})
    else:
        print('Need to organize the data accordingly.')
        train_dataset = {'images': [], 'annotations': [], 'categories': []}
        val_dataset = {'images': [], 'annotations': [], 'categories': []}
        test_dataset = {'images': [], 'annotations': [], 'categories': []}

        # category id starts from 0
        for i, cls in enumerate(classes, 0):
            train_dataset['categories'].append({'id': i, 'name': cls})
            val_dataset['categories'].append({'id': i, 'name': cls})
            test_dataset['categories'].append({'id': i, 'name': cls})
        train_img, val_img, test_img = organize_by_existing_files(
            image_dir, existed_categories)

    obj_count = 0
    skipped = 0
    converted = 0
    for idx, image in enumerate(mmengine.track_iter_progress(indices)):
        img_info_dict, image_height, image_width = get_image_info(
            yolo_image_dir, idx, image)

        if existed_categories != []:
            if image in train_img:
                dataset = train_dataset
            elif image in val_img:
                dataset = val_dataset
            elif image in test_img:
                dataset = test_dataset

        dataset['images'].append(img_info_dict)

        img_name = osp.splitext(image)[0]
        label_path = f'{osp.join(yolo_label_dir, img_name)}.txt'
        if not osp.exists(label_path):
            # if current image is not annotated or the annotation file failed
            print(
                f'WARNING: {label_path} does not exist. Please check the file.'
            )
            skipped += 1
            continue

        with open(label_path) as f:
            labels = f.readlines()
            for label in labels:
                coco_info, obj_count = convert_bbox_info(
                    label, idx, obj_count, image_height, image_width)
                dataset['annotations'].append(coco_info)
        converted += 1

    # saving results to result json
    if existed_categories == []:
        out_file = osp.join(image_dir, 'annotations/result.json')
        print(f'Saving converted results to {out_file} ...')
        mmengine.dump(dataset, out_file)
    else:
        for category in existed_categories:
            out_file = osp.join(output_folder, f'{category}.json')
            print(f'Saving converted results to {out_file} ...')
            if category == 'train':
                mmengine.dump(train_dataset, out_file)
            elif category == 'val':
                mmengine.dump(val_dataset, out_file)
            elif category == 'test':
                mmengine.dump(test_dataset, out_file)

    # simple statistics
    print(f'Process finished! Please check at {output_folder} .')
    print(f'Number of images found: {total}, converted: {converted},',
          f'and skipped: {skipped}. Total annotation count: {obj_count}.')
    print('You can use tools/analysis_tools/browse_coco_json.py to visualize!')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'image_dir',
        type=str,
        help='dataset directory with ./images and ./labels, classes.txt, etc.')
    arg = parser.parse_args()
    convert_yolo_to_coco(arg.image_dir)