RT-ODLab/dataset/scripts/data_to_h5py.py

import cv2
import h5py
import os
import argparse
import numpy as np
import sys

sys.path.append('..')
from voc import VOCDetection
from coco import COCODataset

# ---------------------- Opt ----------------------
parser = argparse.ArgumentParser(description='Cache-Dataset')
parser.add_argument('-d', '--dataset', default='voc',
                    help='coco, voc, widerface, crowdhuman')
parser.add_argument('--root', default='/Users/liuhaoran/Desktop/python_work/object-detection/dataset/',
                    help='data root')
parser.add_argument('-size', '--img_size', default=640, type=int,
                    help='input image size.')
parser.add_argument('--mosaic', default=None, type=float,
                    help='mosaic augmentation.')
parser.add_argument('--mixup', default=None, type=float,
                    help='mixup augmentation.')
parser.add_argument('--keep_ratio', action="store_true", default=False,
                    help='keep aspect ratio.')

args = parser.parse_args()


# ---------------------- Build Dataset ----------------------
if args.dataset == 'voc':
    root = os.path.join(args.root, 'VOCdevkit')
    dataset = VOCDetection(args.img_size, root)
elif args.dataset == 'coco':
    root = os.path.join(args.root, 'COCO')
    dataset = COCODataset(args.img_size, args.root)
print('Data length: ', len(dataset))


# ---------------------- Main Process ----------------------
for i in range(len(dataset)):
    # load an image
    image, image_id = dataset.pull_image(i)
    orig_h, orig_w, _ = image.shape

    # resize image
    if args.keep_ratio:
        r = args.img_size / max(orig_h, orig_w)
        if r != 1: 
            interp = cv2.INTER_LINEAR
            new_size = (int(orig_w * r), int(orig_h * r))
            image = cv2.resize(image, new_size, interpolation=interp)
        pad_image = np.ones([args.img_size, args.img_size, image.shape[2]], dtype=image.dtype) * 114
        pad_image[:new_size[1], :new_size[0]] = image
        image = pad_image
    else:
        image = cv2.resize(image, (int(args.img_size), int(args.img_size)))

    cv2.imshow('image', image)
    # cv2.imwrite(str(i)+'.jpg', img)
    cv2.waitKey(0)