From c948e498abc85ba0c282fc8783b0111d0bfbfeda Mon Sep 17 00:00:00 2001
From: yaoxuefeng <yaoxuefeng@baidu.com>
Date: Thu, 7 May 2020 23:12:14 +0800
Subject: [PATCH] add tools for dataset download

---
 tools/tools.py | 133 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 133 insertions(+)
 create mode 100644 tools/tools.py

diff --git a/tools/tools.py b/tools/tools.py
new file mode 100644
index 00000000..da34a027
--- /dev/null
+++ b/tools/tools.py
@@ -0,0 +1,133 @@
+import os
+import time
+import shutil
+import requests
+import sys
+import tarfile
+import zipfile
+import platform
+import functools
+
+lasttime = time.time()
+FLUSH_INTERVAL = 0.1
+
+LOCAL_PATH = os.path.dirname(os.path.abspath(__file__))
+
+
+def get_platform():
+    return platform.platform()
+
+
+def is_windows():
+    return get_platform().lower().startswith("windows")
+
+
+def progress(str, end=False):
+    global lasttime
+    if end:
+        str += "\n"
+        lasttime = 0
+    if time.time() - lasttime >= FLUSH_INTERVAL:
+        sys.stdout.write("\r%s" % str)
+        lasttime = time.time()
+        sys.stdout.flush()
+
+
+def download_file(url, savepath, print_progress):
+    r = requests.get(url, stream=True)
+    total_length = r.headers.get('content-length')
+
+    if total_length is None:
+        with open(savepath, 'wb') as f:
+            shutil.copyfileobj(r.raw, f)
+    else:
+        with open(savepath, 'wb') as f:
+            dl = 0
+            total_length = int(total_length)
+            starttime = time.time()
+            if print_progress:
+                print("Downloading %s" % os.path.basename(savepath))
+            for data in r.iter_content(chunk_size=4096):
+                dl += len(data)
+                f.write(data)
+                if print_progress:
+                    done = int(50 * dl / total_length)
+                    progress("[%-50s] %.2f%%" %
+                             ('=' * done, float(100 * dl) / total_length))
+        if print_progress:
+            progress("[%-50s] %.2f%%" % ('=' * 50, 100), end=True)
+
+
+def _uncompress_file(filepath, extrapath, delete_file, print_progress):
+    if print_progress:
+        print("Uncompress %s" % os.path.basename(filepath))
+
+    if filepath.endswith("zip"):
+        handler = _uncompress_file_zip
+    elif filepath.endswith("tgz"):
+        handler = _uncompress_file_tar
+    else:
+        handler = functools.partial(_uncompress_file_tar, mode="r")
+
+    for total_num, index, rootpath in handler(filepath, extrapath):
+        if print_progress:
+            done = int(50 * float(index) / total_num)
+            progress("[%-50s] %.2f%%" %
+                     ('=' * done, float(100 * index) / total_num))
+    if print_progress:
+        progress("[%-50s] %.2f%%" % ('=' * 50, 100), end=True)
+
+    if delete_file:
+        os.remove(filepath)
+
+    return rootpath
+
+
+def _uncompress_file_zip(filepath, extrapath):
+    files = zipfile.ZipFile(filepath, 'r')
+    filelist = files.namelist()
+    rootpath = filelist[0]
+    total_num = len(filelist)
+    for index, file in enumerate(filelist):
+        files.extract(file, extrapath)
+        yield total_num, index, rootpath
+    files.close()
+    yield total_num, index, rootpath
+
+
+def _uncompress_file_tar(filepath, extrapath, mode="r:gz"):
+    files = tarfile.open(filepath, mode)
+    filelist = files.getnames()
+    total_num = len(filelist)
+    rootpath = filelist[0]
+    for index, file in enumerate(filelist):
+        files.extract(file, extrapath)
+        yield total_num, index, rootpath
+    files.close()
+    yield total_num, index, rootpath
+
+
+def download_file_and_uncompress(url,
+                                 savepath=None,
+                                 savename=None,
+                                 extrapath=None,
+                                 print_progress=True,
+                                 cover=False,
+                                 delete_file=False):
+    if savepath is None:
+        savepath = "."
+
+    if extrapath is None:
+        extrapath = "."
+
+    if savename is None:
+        savename = url.split("/")[-1]
+    savepath = os.path.join(savepath, savename)
+
+    if cover:
+        if os.path.exists(savepath):
+            shutil.rmtree(savepath)
+
+    if not os.path.exists(savepath):
+        download_file(url, savepath, print_progress)
+    _ = _uncompress_file(savepath, extrapath, delete_file, print_progress)
-- 
GitLab