yolov6/data/vis_dataset.py

# coding=utf-8
# Description:  visualize yolo label image.

import argparse
import os
import cv2
import numpy as np

IMG_FORMATS = ["bmp", "jpg", "jpeg", "png", "tif", "tiff", "dng", "webp", "mpo"]


def main(args):
    img_dir, label_dir, class_names = args.img_dir, args.label_dir, args.class_names

    label_map = dict()
    for class_id, classname in enumerate(class_names):
        label_map[class_id] = classname

    for file in os.listdir(img_dir):
        if file.split('.')[-1] not in IMG_FORMATS:
            print(f'[Warning]: Non-image file {file}')
            continue
        img_path = os.path.join(img_dir, file)
        label_path = os.path.join(label_dir, file[: file.rindex('.')] + '.txt')

        try:
            img_data = cv2.imread(img_path)
            height, width, _ = img_data.shape
            color = [tuple(np.random.choice(range(256), size=3)) for i in class_names]
            thickness = 2

            with open(label_path, 'r') as f:
                for bbox in f:
                    cls, x_c, y_c, w, h = [float(v) if i > 0 else int(v) for i, v in enumerate(bbox.split('\n')[0].split(' '))]

                    x_tl = int((x_c - w / 2) * width)
                    y_tl = int((y_c - h / 2) * height)
                    cv2.rectangle(img_data, (x_tl, y_tl), (x_tl + int(w * width), y_tl + int(h * height)), tuple([int(x) for x in color[cls]]), thickness)
                    cv2.putText(img_data, label_map[cls], (x_tl, y_tl - 10), cv2.FONT_HERSHEY_COMPLEX, 1, tuple([int(x) for x in color[cls]]), thickness)

            cv2.imshow('image', img_data)
            cv2.waitKey(0)
        except Exception as e:
            print(f'[Error]: {e} {img_path}')
    print('======All Done!======')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', default='VOCdevkit/voc_07_12/images')
    parser.add_argument('--label_dir', default='VOCdevkit/voc_07_12/labels')
    parser.add_argument('--class_names', default=['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
                        'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'])

    args = parser.parse_args()
    print(args)

    main(args)