Spaces:
Sleeping
Sleeping
File size: 3,569 Bytes
2232b2c |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
from __future__ import print_function
from pathlib import Path
import tarfile
import requests
from warnings import warn
from zipfile import ZipFile
from bs4 import BeautifulSoup
class GetData(object):
"""A Python script for downloading CycleGAN or pix2pix datasets.
Parameters:
technique (str) -- One of: 'cyclegan' or 'pix2pix'.
verbose (bool) -- If True, print additional information.
Examples:
>>> from util.get_data import GetData
>>> gd = GetData(technique='cyclegan')
>>> new_data_path = gd.get(save_path='./datasets') # options will be displayed.
Alternatively, You can use bash scripts: 'scripts/download_pix2pix_model.sh'
and 'scripts/download_cyclegan_model.sh'.
"""
def __init__(self, technique="cyclegan", verbose=True):
url_dict = {
"pix2pix": "http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/",
"cyclegan": "http://efrosgans.eecs.berkeley.edu/pix2pix/datasets",
}
self.url = url_dict.get(technique.lower())
self._verbose = verbose
def _print(self, text):
if self._verbose:
print(text)
@staticmethod
def _get_options(r):
soup = BeautifulSoup(r.text, "lxml")
options = [h.text for h in soup.find_all("a", href=True) if h.text.endswith((".zip", "tar.gz"))]
return options
def _present_options(self):
r = requests.get(self.url)
options = self._get_options(r)
print("Options:\n")
for i, o in enumerate(options):
print("{0}: {1}".format(i, o))
choice = input("\nPlease enter the number of the " "dataset above you wish to download:")
return options[int(choice)]
def _download_data(self, dataset_url, save_path):
save_path = Path(save_path)
if not save_path.is_dir():
save_path.mkdir(parents=True, exist_ok=True)
base = Path(dataset_url).name
temp_save_path = save_path / base
with open(temp_save_path, "wb") as f:
r = requests.get(dataset_url)
f.write(r.content)
if base.endswith(".tar.gz"):
obj = tarfile.open(temp_save_path)
elif base.endswith(".zip"):
obj = ZipFile(temp_save_path, "r")
else:
raise ValueError("Unknown File Type: {0}.".format(base))
self._print("Unpacking Data...")
obj.extractall(save_path)
obj.close()
temp_save_path.unlink()
def get(self, save_path, dataset=None):
"""
Download a dataset.
Parameters:
save_path (str) -- A directory to save the data to.
dataset (str) -- (optional). A specific dataset to download.
Note: this must include the file extension.
If None, options will be presented for you
to choose from.
Returns:
save_path_full (str) -- the absolute path to the downloaded data.
"""
if dataset is None:
selected_dataset = self._present_options()
else:
selected_dataset = dataset
save_path_full = Path(save_path) / selected_dataset.split(".")[0]
if save_path_full.is_dir():
warn(f"\n'{save_path_full}' already exists. Voiding Download.")
else:
self._print("Downloading Data...")
url = f"{self.url}/{selected_dataset}"
self._download_data(url, save_path=save_path)
return save_path_full.resolve()
|