download.py 13.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import os.path as osp
import shutil
import requests
import tqdm
import hashlib
25 26
import binascii
import base64
27 28 29
import tarfile
import zipfile

K
Kaipeng Deng 已提交
30
from .voc_utils import create_list
31 32 33 34

import logging
logger = logging.getLogger(__name__)

35 36 37 38
__all__ = [
    'get_weights_path', 'get_dataset_path', 'download_dataset',
    'create_voc_list'
]
39 40 41 42

WEIGHTS_HOME = osp.expanduser("~/.cache/paddle/weights")
DATASET_HOME = osp.expanduser("~/.cache/paddle/dataset")

K
Kaipeng Deng 已提交
43
# dict of {dataset_name: (download_info, sub_dirs)}
K
Kaipeng Deng 已提交
44
# download info: [(url, md5sum)]
45 46
DATASETS = {
    'coco': ([
W
wangguanzhong 已提交
47 48 49 50 51 52 53 54 55
        (
            'http://images.cocodataset.org/zips/train2017.zip',
            'cced6f7f71b7629ddf16f17bbcfab6b2', ),
        (
            'http://images.cocodataset.org/zips/val2017.zip',
            '442b8da7639aecaf257c1dceb8ba8c80', ),
        (
            'http://images.cocodataset.org/annotations/annotations_trainval2017.zip',
            'f4bbac642086de4f52a3fdda2de5fa2c', ),
56 57
    ], ["annotations", "train2017", "val2017"]),
    'voc': ([
W
wangguanzhong 已提交
58 59 60 61 62 63 64 65 66
        (
            'http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar',
            '6cd6e144f989b92b3379bac3b3de84fd', ),
        (
            'http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar',
            'c52e279531787c972589f7e41ab4ae64', ),
        (
            'http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar',
            'b6e924de25625d8de591ea690078ad9f', ),
K
Kaipeng Deng 已提交
67
    ], ["VOCdevkit/VOC2012", "VOCdevkit/VOC2007"]),
G
Guanghua Yu 已提交
68 69 70 71 72 73 74 75 76 77 78
    'wider_face': ([
        (
            'https://dataset.bj.bcebos.com/wider_face/WIDER_train.zip',
            '3fedf70df600953d25982bcd13d91ba2', ),
        (
            'https://dataset.bj.bcebos.com/wider_face/WIDER_val.zip',
            'dfa7d7e790efa35df3788964cf0bbaea', ),
        (
            'https://dataset.bj.bcebos.com/wider_face/wider_face_split.zip',
            'a4a898d6193db4b9ef3260a68bad0dc7', ),
    ], ["WIDER_train", "WIDER_val", "wider_face_split"]),
W
wangguanzhong 已提交
79
    'fruit': ([(
80 81 82
        'https://dataset.bj.bcebos.com/PaddleDetection_demo/fruit.tar',
        'baa8806617a54ccf3685fa7153388ae6', ), ],
              ['Annotations', 'JPEGImages']),
C
cnn 已提交
83 84 85 86 87 88
    'roadsign_voc': ([(
        'https://paddlemodels.bj.bcebos.com/object_detection/roadsign_voc.tar',
        '8d629c0f880dd8b48de9aeff44bf1f3e', ), ], ['annotations', 'images']),
    'roadsign_coco': ([(
        'https://paddlemodels.bj.bcebos.com/object_detection/roadsign_coco.tar',
        '49ce5a9b5ad0d6266163cd01de4b018e', ), ], ['annotations', 'images']),
W
wangguanzhong 已提交
89
    'objects365': (),
90 91 92 93 94 95 96 97 98
}

DOWNLOAD_RETRY_LIMIT = 3


def get_weights_path(url):
    """Get weights path from WEIGHT_HOME, if not exists,
    download it from url.
    """
K
Kaipeng Deng 已提交
99 100
    path, _ = get_path(url, WEIGHTS_HOME)
    return path
101 102


103
def get_dataset_path(path, annotation, image_dir):
104 105 106 107 108
    """
    If path exists, return path.
    Otherwise, get dataset path from DATASET_HOME, if not exists,
    download it.
    """
109
    if _dataset_exists(path, annotation, image_dir):
110 111
        return path

K
Kaipeng Deng 已提交
112
    logger.info("Dataset {} is not valid for reason above, try searching {} or "
113 114 115
                "downloading dataset...".format(
                    osp.realpath(path), DATASET_HOME))

116
    data_name = os.path.split(path.strip().lower())[-1]
117
    for name, dataset in DATASETS.items():
118
        if data_name == name:
Y
Yang Zhang 已提交
119 120
            logger.debug("Parse dataset_dir {} as dataset "
                         "{}".format(path, name))
W
wangguanzhong 已提交
121 122
            if name == 'objects365':
                raise NotImplementedError(
123 124 125
                    "Dataset {} is not valid for download automatically. "
                    "Please apply and download the dataset from "
                    "https://www.objects365.org/download.html".format(name))
126
            data_dir = osp.join(DATASET_HOME, name)
K
Kaipeng Deng 已提交
127
            # For voc, only check dir VOCdevkit/VOC2012, VOCdevkit/VOC2007
C
cnn 已提交
128
            if name == 'voc' or name == 'fruit' or name == 'roadsign_voc':
K
Kaipeng Deng 已提交
129 130 131 132 133 134 135 136
                exists = True
                for sub_dir in dataset[1]:
                    check_dir = osp.join(data_dir, sub_dir)
                    if osp.exists(check_dir):
                        logger.info("Found {}".format(check_dir))
                    else:
                        exists = False
                if exists:
137 138
                    return data_dir

K
Kaipeng Deng 已提交
139
            # voc exist is checked above, voc is not exist here
C
cnn 已提交
140
            check_exist = name != 'voc' and name != 'fruit' and name != 'roadsign_voc'
141
            for url, md5sum in dataset[0]:
K
Kaipeng Deng 已提交
142
                get_path(url, data_dir, md5sum, check_exist)
143

K
Kaipeng Deng 已提交
144
            # voc should create list after download
145
            if name == 'voc':
K
Kaipeng Deng 已提交
146
                create_voc_list(data_dir)
147 148 149
            return data_dir

    # not match any dataset in DATASETS
C
cnn 已提交
150 151 152 153 154
    raise ValueError(
        "Dataset {} is not valid and cannot parse dataset type "
        "'{}' for automaticly downloading, which only supports "
        "'voc' , 'coco', 'wider_face', 'fruit' and 'roadsign_voc' currently".
        format(path, osp.split(path)[-1]))
155 156


K
Kaipeng Deng 已提交
157
def create_voc_list(data_dir, devkit_subdir='VOCdevkit'):
Y
Yang Zhang 已提交
158
    logger.debug("Create voc file list...")
K
Kaipeng Deng 已提交
159
    devkit_dir = osp.join(data_dir, devkit_subdir)
160
    year_dirs = [osp.join(devkit_dir, x) for x in os.listdir(devkit_dir)]
K
Kaipeng Deng 已提交
161

K
Kaipeng Deng 已提交
162 163 164
    # NOTE: since using auto download VOC
    # dataset, VOC default label list should be used, 
    # do not generate label_list.txt here. For default
165
    # label, see ../data/source/voc.py
166
    create_list(year_dirs, data_dir)
Y
Yang Zhang 已提交
167
    logger.debug("Create voc file list finished")
K
Kaipeng Deng 已提交
168 169


170 171
def map_path(url, root_dir):
    # parse path after download to decompress under root_dir
172
    fname = osp.split(url)[-1]
173 174 175 176 177 178 179
    zip_formats = ['.zip', '.tar', '.gz']
    fpath = fname
    for zip_format in zip_formats:
        fpath = fpath.replace(zip_format, '')
    return osp.join(root_dir, fpath)


K
Kaipeng Deng 已提交
180
def get_path(url, root_dir, md5sum=None, check_exist=True):
181 182 183 184 185 186 187 188 189 190 191
    """ Download from given url to root_dir.
    if file or directory specified by url is exists under
    root_dir, return the path directly, otherwise download
    from url and decompress it, return the path.

    url (str): download url
    root_dir (str): root dir for downloading, it should be
                    WEIGHTS_HOME or DATASET_HOME
    md5sum (str): md5 sum of download package
    """
    # parse path after download to decompress under root_dir
192
    fullpath = map_path(url, root_dir)
193 194 195 196

    # For same zip file, decompressed directory name different
    # from zip file name, rename by following map
    decompress_name_map = {
K
Kaipeng Deng 已提交
197 198 199
        "VOCtrainval_11-May-2012": "VOCdevkit/VOC2012",
        "VOCtrainval_06-Nov-2007": "VOCdevkit/VOC2007",
        "VOCtest_06-Nov-2007": "VOCdevkit/VOC2007",
200 201 202 203
        "annotations_trainval": "annotations"
    }
    for k, v in decompress_name_map.items():
        if fullpath.find(k) >= 0:
204
            fullpath = osp.join(osp.split(fullpath)[0], v)
205

K
Kaipeng Deng 已提交
206 207 208
    exist_flag = False
    if osp.exists(fullpath) and check_exist:
        exist_flag = True
Y
Yang Zhang 已提交
209
        logger.debug("Found {}".format(fullpath))
210
    else:
K
Kaipeng Deng 已提交
211
        exist_flag = False
212
        fullname = _download(url, root_dir, md5sum)
K
Kaipeng Deng 已提交
213 214 215 216 217

        # new weights format which postfix is 'pdparams' not
        # need to decompress
        if osp.splitext(fullname)[-1] != '.pdparams':
            _decompress(fullname)
218

K
Kaipeng Deng 已提交
219
    return fullpath, exist_flag
220 221


K
Kaipeng Deng 已提交
222 223 224 225 226 227 228
def download_dataset(path, dataset=None):
    if dataset not in DATASETS.keys():
        logger.error("Unknown dataset {}, it should be "
                     "{}".format(dataset, DATASETS.keys()))
        return
    dataset_info = DATASETS[dataset][0]
    for info in dataset_info:
K
Kaipeng Deng 已提交
229
        get_path(info[0], path, info[1], False)
Y
Yang Zhang 已提交
230
    logger.debug("Download dataset {} finished.".format(dataset))
K
Kaipeng Deng 已提交
231 232


233
def _dataset_exists(path, annotation, image_dir):
234 235 236 237
    """
    Check if user define dataset exists
    """
    if not osp.exists(path):
Y
Yang Zhang 已提交
238 239
        logger.debug("Config dataset_dir {} is not exits, "
                     "dataset config is not valid".format(path))
240 241
        return False

242 243
    if annotation:
        annotation_path = osp.join(path, annotation)
C
cnn 已提交
244 245 246
        if not osp.exists(annotation_path):
            logger.error("Config dataset_dir {} is not exits!".format(path))

247
        if not osp.isfile(annotation_path):
C
cnn 已提交
248 249 250
            logger.warning("Config annotation {} is not a "
                           "file, dataset config is not "
                           "valid".format(annotation_path))
251 252 253
            return False
    if image_dir:
        image_path = osp.join(path, image_dir)
C
cnn 已提交
254 255 256
        if not osp.exists(image_path):
            logger.warning("Config dataset_dir {} is not exits!".format(path))

257
        if not osp.isdir(image_path):
Y
Yang Zhang 已提交
258 259 260
            logger.warning("Config image_dir {} is not a "
                           "directory, dataset config is not "
                           "valid".format(image_path))
261
            return False
262 263 264 265 266 267 268 269 270 271 272 273 274
    return True


def _download(url, path, md5sum=None):
    """
    Download from url, save to path.

    url (str): download url
    path (str): download to given path
    """
    if not osp.exists(path):
        os.makedirs(path)

275
    fname = osp.split(url)[-1]
276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292
    fullname = osp.join(path, fname)
    retry_cnt = 0

    while not (osp.exists(fullname) and _md5check(fullname, md5sum)):
        if retry_cnt < DOWNLOAD_RETRY_LIMIT:
            retry_cnt += 1
        else:
            raise RuntimeError("Download from {} failed. "
                               "Retry limit reached".format(url))

        logger.info("Downloading {} from {}".format(fname, url))

        req = requests.get(url, stream=True)
        if req.status_code != 200:
            raise RuntimeError("Downloading from {} failed with code "
                               "{}!".format(url, req.status_code))

K
Kaipeng Deng 已提交
293 294 295 296
        # For protecting download interupted, download to
        # tmp_fullname firstly, move tmp_fullname to fullname
        # after download finished
        tmp_fullname = fullname + "_tmp"
297
        total_size = req.headers.get('content-length')
K
Kaipeng Deng 已提交
298
        with open(tmp_fullname, 'wb') as f:
299 300 301 302 303 304 305 306 307 308 309
            if total_size:
                for chunk in tqdm.tqdm(
                        req.iter_content(chunk_size=1024),
                        total=(int(total_size) + 1023) // 1024,
                        unit='KB'):
                    f.write(chunk)
            else:
                for chunk in req.iter_content(chunk_size=1024):
                    if chunk:
                        f.write(chunk)

310 311 312 313 314 315 316 317 318 319 320 321 322
        # check md5 after download in Content-MD5 in req.headers
        content_md5 = req.headers.get('content-md5')
        if not content_md5 or _md5check(
                tmp_fullname,
                binascii.hexlify(base64.b64decode(content_md5.strip(
                    '"'))).decode()):
            shutil.move(tmp_fullname, fullname)
            return fullname
        else:
            logger.warn(
                "Download from url imcomplete, try downloading again...")
            os.remove(tmp_fullname)
            continue
323 324 325 326 327 328


def _md5check(fullname, md5sum=None):
    if md5sum is None:
        return True

Y
Yang Zhang 已提交
329
    logger.debug("File {} md5 checking...".format(fullname))
330 331 332 333 334 335 336
    md5 = hashlib.md5()
    with open(fullname, 'rb') as f:
        for chunk in iter(lambda: f.read(4096), b""):
            md5.update(chunk)
    calc_md5sum = md5.hexdigest()

    if calc_md5sum != md5sum:
Y
Yang Zhang 已提交
337 338
        logger.warning("File {} md5 check failed, {}(calc) != "
                       "{}(base)".format(fullname, calc_md5sum, md5sum))
339 340 341 342 343 344 345 346 347 348 349 350 351 352
        return False
    return True


def _decompress(fname):
    """
    Decompress for zip and tar file
    """
    logger.info("Decompressing {}...".format(fname))

    # For protecting decompressing interupted,
    # decompress to fpath_tmp directory firstly, if decompress
    # successed, move decompress files to fpath and delete
    # fpath_tmp and remove download compress file.
353
    fpath = osp.split(fname)[0]
354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378
    fpath_tmp = osp.join(fpath, 'tmp')
    if osp.isdir(fpath_tmp):
        shutil.rmtree(fpath_tmp)
        os.makedirs(fpath_tmp)

    if fname.find('tar') >= 0:
        with tarfile.open(fname) as tf:
            tf.extractall(path=fpath_tmp)
    elif fname.find('zip') >= 0:
        with zipfile.ZipFile(fname) as zf:
            zf.extractall(path=fpath_tmp)
    else:
        raise TypeError("Unsupport compress file type {}".format(fname))

    for f in os.listdir(fpath_tmp):
        src_dir = osp.join(fpath_tmp, f)
        dst_dir = osp.join(fpath, f)
        _move_and_merge_tree(src_dir, dst_dir)

    shutil.rmtree(fpath_tmp)
    os.remove(fname)


def _move_and_merge_tree(src, dst):
    """
G
Guanghua Yu 已提交
379
    Move src directory to dst, if dst is already exists,
380 381 382 383
    merge src to dst
    """
    if not osp.exists(dst):
        shutil.move(src, dst)
384 385
    elif osp.isfile(src):
        shutil.move(src, dst)
386 387 388 389 390 391 392 393 394 395 396 397
    else:
        for fp in os.listdir(src):
            src_fp = osp.join(src, fp)
            dst_fp = osp.join(dst, fp)
            if osp.isdir(src_fp):
                if osp.isdir(dst_fp):
                    _move_and_merge_tree(src_fp, dst_fp)
                else:
                    shutil.move(src_fp, dst_fp)
            elif osp.isfile(src_fp) and \
                    not osp.isfile(dst_fp):
                shutil.move(src_fp, dst_fp)