| | 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() |
| |
|
