crashedice/signify

1from __future__ import print_function2import os3import tarfile4import requests5from warnings import warn6from zipfile import ZipFile7from bs4 import BeautifulSoup8from os.path import abspath, isdir, join, basename9 10 11class GetData(object):12    """A Python script for downloading CycleGAN or pix2pix datasets.13 14    Parameters:15        technique (str) -- One of: 'cyclegan' or 'pix2pix'.16        verbose (bool)  -- If True, print additional information.17 18    Examples:19        >>> from util.get_data import GetData20        >>> gd = GetData(technique='cyclegan')21        >>> new_data_path = gd.get(save_path='./datasets')  # options will be displayed.22 23    Alternatively, You can use bash scripts: 'scripts/download_pix2pix_model.sh'24    and 'scripts/download_cyclegan_model.sh'.25    """26 27    def __init__(self, technique='cyclegan', verbose=True):28        url_dict = {29            'pix2pix': 'http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/',30            'cyclegan': 'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets'31        }32        self.url = url_dict.get(technique.lower())33        self._verbose = verbose34 35    def _print(self, text):36        if self._verbose:37            print(text)38 39    @staticmethod40    def _get_options(r):41        soup = BeautifulSoup(r.text, 'lxml')42        options = [h.text for h in soup.find_all('a', href=True)43                   if h.text.endswith(('.zip', 'tar.gz'))]44        return options45 46    def _present_options(self):47        r = requests.get(self.url)48        options = self._get_options(r)49        print('Options:\n')50        for i, o in enumerate(options):51            print("{0}: {1}".format(i, o))52        choice = input("\nPlease enter the number of the "53                       "dataset above you wish to download:")54        return options[int(choice)]55 56    def _download_data(self, dataset_url, save_path):57        if not isdir(save_path):58            os.makedirs(save_path)59 60        base = basename(dataset_url)61        temp_save_path = join(save_path, base)62 63        with open(temp_save_path, "wb") as f:64            r = requests.get(dataset_url)65            f.write(r.content)66 67        if base.endswith('.tar.gz'):68            obj = tarfile.open(temp_save_path)69        elif base.endswith('.zip'):70            obj = ZipFile(temp_save_path, 'r')71        else:72            raise ValueError("Unknown File Type: {0}.".format(base))73 74        self._print("Unpacking Data...")75        obj.extractall(save_path)76        obj.close()77        os.remove(temp_save_path)78 79    def get(self, save_path, dataset=None):80        """81 82        Download a dataset.83 84        Parameters:85            save_path (str) -- A directory to save the data to.86            dataset (str)   -- (optional). A specific dataset to download.87                            Note: this must include the file extension.88                            If None, options will be presented for you89                            to choose from.90 91        Returns:92            save_path_full (str) -- the absolute path to the downloaded data.93 94        """95        if dataset is None:96            selected_dataset = self._present_options()97        else:98            selected_dataset = dataset99 100        save_path_full = join(save_path, selected_dataset.split('.')[0])101 102        if isdir(save_path_full):103            warn("\n'{0}' already exists. Voiding Download.".format(104                save_path_full))105        else:106            self._print('Downloading Data...')107            url = "{0}/{1}".format(self.url, selected_dataset)108            self._download_data(url, save_path=save_path)109 110        return abspath(save_path_full)111