forked from zhangbin0917/PyTorch-Encoding
-
Notifications
You must be signed in to change notification settings - Fork 0
/
prepare_coco.py
68 lines (61 loc) · 2.61 KB
/
prepare_coco.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
"""Prepare MS COCO datasets"""
import os
import shutil
import argparse
import zipfile
from encoding.utils import download, mkdir
_TARGET_DIR = os.path.expanduser('~/.encoding/data')
def parse_args():
parser = argparse.ArgumentParser(
description='Initialize MS COCO dataset.',
epilog='Example: python mscoco.py --download-dir ~/mscoco',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--download-dir', type=str, default=None, help='dataset directory on disk')
args = parser.parse_args()
return args
def download_coco(path, overwrite=False):
_DOWNLOAD_URLS = [
('http://images.cocodataset.org/zips/train2017.zip',
'10ad623668ab00c62c096f0ed636d6aff41faca5'),
('http://images.cocodataset.org/zips/val2017.zip',
'4950dc9d00dbe1c933ee0170f5797584351d2a41'),
('http://images.cocodataset.org/annotations/annotations_trainval2017.zip',
'8551ee4bb5860311e79dace7e79cb91e432e78b3'),
#('http://images.cocodataset.org/annotations/stuff_annotations_trainval2017.zip',
# '46cdcf715b6b4f67e980b529534e79c2edffe084'),
#('http://images.cocodataset.org/zips/test2017.zip',
# '99813c02442f3c112d491ea6f30cecf421d0e6b3'),
('https://hangzh.s3.amazonaws.com/encoding/data/coco/train_ids.pth',
'12cd266f97c8d9ea86e15a11f11bcb5faba700b6'),
('https://hangzh.s3.amazonaws.com/encoding/data/coco/val_ids.pth',
'4ce037ac33cbf3712fd93280a1c5e92dae3136bb'),
]
mkdir(path)
for url, checksum in _DOWNLOAD_URLS:
filename = download(url, path=path, overwrite=overwrite, sha1_hash=checksum)
# extract
if os.path.splitext(filename)[1] == '.zip':
with zipfile.ZipFile(filename) as zf:
zf.extractall(path=path)
else:
shutil.move(filename, os.path.join(path, 'annotations/'+os.path.basename(filename)))
def install_coco_api():
repo_url = "https://github.com/cocodataset/cocoapi"
os.system("git clone " + repo_url)
os.system("cd cocoapi/PythonAPI/ && python setup.py install")
shutil.rmtree('cocoapi')
try:
import pycocotools
except Exception:
print("Installing COCO API failed, please install it manually %s"%(repo_url))
if __name__ == '__main__':
args = parse_args()
mkdir(os.path.expanduser('~/.encoding/data'))
if args.download_dir is not None:
if os.path.isdir(_TARGET_DIR):
os.remove(_TARGET_DIR)
# make symlink
os.symlink(args.download_dir, _TARGET_DIR)
else:
download_coco(_TARGET_DIR, overwrite=False)
install_coco_api()