| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061 |
- 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 ----------------------
- image_dict = dict()
- 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)
- else:
- image = cv2.resize(image, (int(args.img_size), int(args.img_size)))
- image_dict[image_id] = image
- cv2.imshow('image', image)
- # cv2.imwrite(str(i)+'.jpg', img)
- cv2.waitKey(0)
|
