Open Images Dataset V4 Train - Valid - Test 解析为 Darknet-YOLO 训练数据 (backpack-handbag-suitcase)

Open Images Dataset V4 Train - Valid - Test 解析为 Darknet-YOLO 训练数据 (backpack-handbag-suitcase)

1. Backpack - Handbag - Suitcase

# category_id = '/m/01940j' - Backpack - 0
# category_id = '/m/080hkjn' - Handbag - 1
# category_id = '/m/01s55n' - Suitcase - 2
# category_id = '/m/0hf58v5' - Luggage and bags - 3

2. open_images_dataset_v4_parser.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Yongqiang Cheng

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import csv
import os
import shutil

import cv2

# category_id = '/m/01940j' - Backpack - 0
# category_id = '/m/080hkjn' - Handbag - 1
# category_id = '/m/01s55n' - Suitcase - 2
# category_id = '/m/0hf58v5' - Luggage and bags - 3

labels = {'/m/080hkjn': [1, 'Handbag'], '/m/01940j': [0, 'Backpack'], '/m/01s55n': [2, 'Suitcase'],
          '/m/0hf58v5': [3, 'Luggage and bags']}

luggage_and_bags = '/m/0hf58v5'

annotation_train = "/media/strong/bdc8242d-bed5-479c-ba62-484e35ce7e76/open_img/CSV/train-name_label_annotation.csv"
annotation_valid = "/media/strong/bdc8242d-bed5-479c-ba62-484e35ce7e76/open_img/CSV/valid_name_label_annotation.csv"
annotation_test = "/media/strong/bdc8242d-bed5-479c-ba62-484e35ce7e76/open_img/CSV/test_name_label_annotation.csv"

image_train = "/media/strong/bdc8242d-bed5-479c-ba62-484e35ce7e76/open_img/train-img"
image_valid = "/media/strong/bdc8242d-bed5-479c-ba62-484e35ce7e76/open_img/valid-img"
image_test = "/media/strong/bdc8242d-bed5-479c-ba62-484e35ce7e76/open_img/test-img"

task = "valid"  # "train" - "valid" - "test"
train_images = "/home/strong/data_update/Open_Images_Dataset_V4/%sv4_train/JPEGImages" % (task)
train_labels = "/home/strong/data_update/Open_Images_Dataset_V4/%sv4_train/labels" % (task)
images_labels_path = "/home/strong/data_update/Open_Images_Dataset_V4/%sv4_train/JPEGImages-labels" % (task)

annotation_v4 = annotation_valid
image_v4 = image_valid

if not os.path.exists(train_images):
    os.makedirs(train_images)

if not os.path.exists(train_labels):
    os.makedirs(train_labels)

if not os.path.exists(images_labels_path):
    os.makedirs(images_labels_path)

img_file_names = os.listdir(image_v4 + "/")


# img_file_names.sort()

# cv2.namedWindow("image_labels", cv2.WINDOW_NORMAL)

# size = (image_width, image_height)
# box = (float(box_xmin_pixel), float(box_xmax_pixel), float(box_ymin_pixel), float(box_ymax_pixel))
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]

    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh

    return (x, y, w, h)


def bbox_fix(box_xmin_pixel, box_xmax_pixel, box_ymin_pixel, box_ymax_pixel, img_file_name):
    if (box_xmin_pixel < 0) or (box_xmin_pixel > (img_width - 1)):
        print("box_xmin_pixel error!")
        print(img_file_name)

    if (box_xmax_pixel < 0) or (box_xmax_pixel > (img_width - 1)):
        print("box_xmax_pixel error!")
        print(img_file_name)

    if (box_ymin_pixel < 0) or (box_ymin_pixel > (img_height - 1)):
        print("box_ymin_pixel error!")
        print(img_file_name)

    if (box_ymax_pixel < 0) or (box_ymax_pixel > (img_height - 1)):
        print("box_ymax_pixel error!")
        print(img_file_name)

    if (box_xmin_pixel < 1):
        box_xmin_pixel = 1
        print("box_xmin_pixel bug fix!")
        print(img_file_name)

    if (box_xmin_pixel > (img_width - 2)):
        box_xmin_pixel = (img_width - 2)
        print("box_xmin_pixel bug fix!")
        print(img_file_name)

    if (box_ymin_pixel < 1):
        box_ymin_pixel = 1
        print("box_ymin_pixel bug fix!")
        print(img_file_name)

    if (box_ymin_pixel > (img_height - 2)):
        box_ymin_pixel = (img_height - 2)
        print("box_ymin_pixel bug fix!")
        print(img_file_name)

    if (box_xmax_pixel > (img_width - 2)):
        box_xmax_pixel = (img_width - 2)
        print("box_xmax_pixel bug fix!")
        print(img_file_name)

    if (box_xmax_pixel < 1):
        box_xmax_pixel = 1
        print("box_xmax_pixel bug fix!")
        print(img_file_name)

    if (box_ymax_pixel > (img_height - 2)):
        box_ymax_pixel = (img_height - 2)
        print("box_ymax_pixel bug fix!")
        print(img_file_name)

    if (box_ymax_pixel <= 1):
        box_ymax_pixel = 1
        print("box_ymax_pixel bug fix!")
        print(img_file_name)

    return box_xmin_pixel, box_xmax_pixel, box_ymin_pixel, box_ymax_pixel


reader_list = []
f = open(annotation_v4, 'r')
reader = csv.reader(f)
# item: ['ImageID', 'Source', 'LabelName', 'Confidence', 'XMin', 'XMax', 'YMin', 'YMax', 'IsOccluded', 'IsTruncated', 'IsGroupOf', 'IsDepiction', 'IsInside']
for item in reader:
    if "ImageID" == item[0]:
        continue
    pass
    # print("item: %s" % (item))

    if item[2] in labels.keys():
        reader_list.append(item)
    pass
pass
f.close()

bag_flag = 0
frame_num = 0
for img_file_name in img_file_names:
    if ".jpg" not in img_file_name:
        continue
    pass

    frame_num += 1
    print("frame_num = %s" % (frame_num))

    bag_flag = 0
    img_file = image_v4 + '/' + img_file_name.strip()

    img_name = img_file_name.strip()

    # item: ['ImageID', 'Source', 'LabelName', 'Confidence', 'XMin', 'XMax', 'YMin', 'YMax', 'IsOccluded', 'IsTruncated', 'IsGroupOf', 'IsDepiction', 'IsInside']
    for item in reader_list:
        # print("item: %s" % (item))
        if item[2] not in labels.keys():
            continue
        pass

        image_id = item[0]
        image_id_jpg = item[0] + '.jpg'
        if str(img_name) != str(image_id_jpg):
            continue
        pass

        # item[10] = 'IsGroupOf'
        if '0' != item[10]:
            bag_flag = 0
            break
        pass

        # item[11] = 'IsDepiction'
        if '0' != item[11]:
            bag_flag = 0
            break
        pass

        # item[9] = 'IsTruncated'
        if '0' != item[9]:
            bag_flag = 0
            break
        pass

        # category_id = '/m/0hf58v5' - Luggage and bags - 3
        if item[2] == luggage_and_bags:
            bag_flag = 0
            break
        pass

        img = cv2.imread(img_file)
        # print(img_file)

        imgcopy = img.copy()
        img_height, img_width, img_channel = img.shape

        x_min = float(item[4])
        x_max = float(item[5])
        y_min = float(item[6])
        y_max = float(item[7])

        # "bbox": [xmin, ymin, width, height]
        box_xmin_pixel = float(x_min) * img_width
        box_xmax_pixel = float(x_max) * img_width
        box_ymin_pixel = float(y_min) * img_height
        box_ymax_pixel = float(y_max) * img_height

        box_xmin_pixel, box_xmax_pixel, box_ymin_pixel, box_ymax_pixel = bbox_fix(box_xmin_pixel,
                                                                                  box_xmax_pixel,
                                                                                  box_ymin_pixel,
                                                                                  box_ymax_pixel,
                                                                                  img_file_name)

        if (box_xmin_pixel < 0) or (box_xmin_pixel > (img_width - 1)):
            print("box_xmin_pixel continue error!")
            print(img_file_name)
            continue

        if (box_xmax_pixel < 0) or (box_xmax_pixel > (img_width - 1)):
            print("box_xmax_pixel continue error!")
            print(img_file_name)
            continue

        if (box_ymin_pixel < 0) or (box_ymin_pixel > (img_height - 1)):
            print("box_ymin_pixel continue error!")
            print(img_file_name)
            continue

        if (box_ymax_pixel < 0) or (box_ymax_pixel > (img_height - 1)):
            print("box_ymax_pixel continue error!")
            print(img_file_name)
            continue

        box_w_pixel = box_xmax_pixel - box_xmin_pixel + 1
        box_h_pixel = box_ymax_pixel - box_ymin_pixel + 1

        if (box_w_pixel < 20) or (box_h_pixel < 20):
            print("box_w_pixel & box_h_pixel continue error!")
            print(img_file_name)
            continue

        bag_flag = 1
        train_label_txt = open("%s/%s" % (train_labels, image_id + '.txt'), 'a+')
        lable = labels[item[2]][0]

        b = (float(box_xmin_pixel), float(box_xmax_pixel), float(box_ymin_pixel), float(box_ymax_pixel))
        bb = convert((img_width, img_height), b)

        train_label_txt.write(str(lable) + " " + " ".join([str(a) for a in bb]) + '\n')
        train_label_txt.close()

        start_x = int((bb[0] - bb[2] / 2.0) * img_width)
        start_y = int((bb[1] - bb[3] / 2.0) * img_height)
        end_x = int((bb[0] + bb[2] / 2.0) * img_width)
        end_y = int((bb[1] + bb[3] / 2.0) * img_height)

        cv2.rectangle(imgcopy, (start_x, start_y), (end_x, end_y), (0, 255, 0), 2)
        cv2.putText(imgcopy, str(lable), (start_x + 2, start_y + 26), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    pass

    if 1 == bag_flag:
        cv2.putText(imgcopy, img_file_name, (16, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255))
        shutil.copy("%s" % (img_file), "%s/" % (train_images))
        # cv2.imwrite("%s/%s" % (train_images, img_file_name.strip()), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
        cv2.imwrite("%s/%s" % (images_labels_path, img_file_name.strip()), imgcopy)

        cv2.imshow("image_labels", imgcopy)
        keyboard = cv2.waitKey(5) & 0xFF
        # # wait for ESC key to exit
        if keyboard == 27:
            break
        pass
    pass
cv2.destroyAllWindows()
pass

if __name__ == '__main__':
    current_directory = os.path.dirname(os.path.abspath(__file__))
    print("current_directory:", current_directory)

3. python3 ./open_images_dataset_v4_parser.py

/usr/bin/python3.5 /home/strong/data_update/Open_Images_Dataset_V4_parser/open_images_dataset_v4_parser.py

来源：CSDN

作者：ForeverStrong

链接：https://blog.csdn.net/chengyq116/article/details/103532275