dmlc--dgl
Tong He
cbee427839
* add working scripts * add frcnn training script * remove redundent files * refactor validation computation, will optimize sgdet and training * validation finally finished * f-rcnn training * test reldn * rm file * update reldn training * data preprocess to h5 * temp * use coco json * fix conflict * new obj dataset for detection * update training * before cleanup * remove abundant files * add arg parse to train * cleanup code file * update * fix * add readme * add ipynb as demo * add demo pic * update readme * add demo script * improve paths * improve readme * add docstrings * fix args description * update readme * add models from s3 * update README Co-authored-by: Minjie Wang <minjie.wang@nyu.edu>
77 行
2.9 KiB
Python
77 行
2.9 KiB
Python
"""Prepare Visual Genome datasets"""
|
|
import os
|
|
import shutil
|
|
import argparse
|
|
import zipfile
|
|
import random
|
|
import json
|
|
import tqdm
|
|
import pickle
|
|
from gluoncv.utils import download, makedirs
|
|
|
|
_TARGET_DIR = os.path.expanduser('~/.mxnet/datasets/visualgenome')
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='Initialize Visual Genome dataset.',
|
|
epilog='Example: python visualgenome.py --download-dir ~/visualgenome',
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument('--download-dir', type=str, default='~/visualgenome/',
|
|
help='dataset directory on disk')
|
|
parser.add_argument('--no-download', action='store_true', help='disable automatic download if set')
|
|
parser.add_argument('--overwrite', action='store_true', help='overwrite downloaded files if set, in case they are corrupted')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
def download_vg(path, overwrite=False):
|
|
_DOWNLOAD_URLS = [
|
|
('https://cs.stanford.edu/people/rak248/VG_100K_2/images.zip',
|
|
'a055367f675dd5476220e9b93e4ca9957b024b94'),
|
|
('https://cs.stanford.edu/people/rak248/VG_100K_2/images2.zip',
|
|
'2add3aab77623549e92b7f15cda0308f50b64ecf'),
|
|
]
|
|
makedirs(path)
|
|
for url, checksum in _DOWNLOAD_URLS:
|
|
filename = download(url, path=path, overwrite=overwrite, sha1_hash=checksum)
|
|
# extract
|
|
if filename.endswith('zip'):
|
|
with zipfile.ZipFile(filename) as zf:
|
|
zf.extractall(path=path)
|
|
# move all images into folder `VG_100K`
|
|
vg_100k_path = os.path.join(path, 'VG_100K')
|
|
vg_100k_2_path = os.path.join(path, 'VG_100K_2')
|
|
files_2 = os.listdir(vg_100k_2_path)
|
|
for fl in files_2:
|
|
shutil.move(os.path.join(vg_100k_2_path, fl),
|
|
os.path.join(vg_100k_path, fl))
|
|
|
|
def download_json(path, overwrite=False):
|
|
url = 'https://data.dgl.ai/dataset/vg.zip'
|
|
output = 'vg.zip'
|
|
download(url, path=path)
|
|
with zipfile.ZipFile(output) as zf:
|
|
zf.extractall(path=path)
|
|
json_path = os.path.join(path, 'vg')
|
|
json_files = os.listdir(json_path)
|
|
for fl in json_files:
|
|
shutil.move(os.path.join(json_path, fl),
|
|
os.path.join(path, fl))
|
|
os.rmdir(json_path)
|
|
|
|
if __name__ == '__main__':
|
|
args = parse_args()
|
|
path = os.path.expanduser(args.download_dir)
|
|
if not os.path.isdir(path):
|
|
if args.no_download:
|
|
raise ValueError(('{} is not a valid directory, make sure it is present.'
|
|
' Or you should not disable "--no-download" to grab it'.format(path)))
|
|
else:
|
|
download_vg(path, overwrite=args.overwrite)
|
|
download_json(path, overwrite=args.overwrite)
|
|
|
|
# make symlink
|
|
makedirs(os.path.expanduser('~/.mxnet/datasets'))
|
|
if os.path.isdir(_TARGET_DIR):
|
|
os.rmdir(_TARGET_DIR)
|
|
os.symlink(path, _TARGET_DIR)
|