MLBench/ReaLens

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