wehub-resource-sync
107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
"""Prepare Visual Genome datasets"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import pickle
|
|
import random
|
|
import shutil
|
|
import zipfile
|
|
|
|
import tqdm
|
|
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)
|