| """Utility functions for providing data directories.""" |
| import os |
| import sys |
| import zipfile |
| import rarfile |
| from urllib.request import FancyURLopener |
| import shutil |
|
|
| import numpy as np |
| import matplotlib.image as mpimg |
|
|
|
|
| class Data(): |
| |
| |
| dirs = None |
|
|
| def __init__(self, data_dir, stat_log_dir, |
| development=True, fast_dir=None): |
| self.development = development |
| self.data_dir = data_dir |
| self.stat_log_dir = stat_log_dir |
| if not os.path.isdir(data_dir): |
| os.makedirs(data_dir) |
|
|
| self._fetch_if_missing() |
|
|
| self.fast_dir = fast_dir |
| if fast_dir: |
| print(">> Copying files to {}".format(fast_dir)) |
| for d in self.dirs: |
| src = os.path.join(data_dir, d) |
| dst = os.path.join(fast_dir, d) |
| if not os.path.isdir(dst): |
| shutil.copytree(src, dst) |
| print(">> Copied {}".format(d)) |
| self.current_dir = fast_dir |
| else: |
| self.current_dir = data_dir |
|
|
| if stat_log_dir: |
| self.stat_log_file = os.path.join(stat_log_dir, |
| self.__class__.__name__ + ".txt") |
| self._ensure_statistics() |
|
|
| def __del__(self): |
| pass |
| |
| |
| |
| |
|
|
| def clear_statistics(self): |
| """Delete saved statistics file if present.""" |
| if self.stat_log_dir and os.path.isfile(self.stat_log_file): |
| os.remove(self.stat_log_file) |
|
|
| def _ensure_statistics(self): |
| """Make sure we know the dataset statistics.""" |
| if os.path.isfile(self.stat_log_file): |
| vals = np.loadtxt(self.stat_log_file) |
| self.mean = vals[0] |
| self.stddev = vals[1] |
| else: |
| print(">> Computing statistics (mean, variance) for {}" |
| .format(self.__class__.__name__)) |
| mean, stddev = self.compute_statistics(self.get_raw_files()) |
| self.mean = mean |
| self.stddev = stddev |
| os.makedirs(self.stat_log_dir, exist_ok=True) |
| np.savetxt(self.stat_log_file, [mean, stddev]) |
| print(">> Statistics complete") |
|
|
| def get_raw_dirs(self): |
| """Should return a list of all dirs containing training images. |
| |
| Note: self.current_dir should be used for loading input data. |
| """ |
| raise NotImplementedError() |
|
|
| def get_raw_files(self): |
| files = [] |
| for d in self.get_raw_dirs(): |
| for path in os.listdir(d): |
| files.append(os.path.join(d, path)) |
| return files |
|
|
| def _fetch_if_missing(self): |
| """A call to this must make subsequent calls to get_raw_files succeed. |
| All subdirs of data_dir listed in self.dirs must exist after this call. |
| """ |
| raise NotImplementedError() |
|
|
| def _download_and_extract(self, url, extract_to, ext='zip'): |
| def _progress(count, block_size, total_size): |
| if total_size > 0: |
| print('\r>> Downloading %s %.1f%%' % (url, |
| float(count * block_size) / float(total_size) * 100.0), end=' ') |
| else: |
| print('\r>> Downloading %s' % (url), end=' ') |
| sys.stdout.flush() |
| urlretrieve = FancyURLopener().retrieve |
| local_zip_path = os.path.join(self.data_dir, 'tmp.' + ext) |
| urlretrieve(url, local_zip_path, _progress) |
| sys.stdout.write("\n>> Finished downloading. Unzipping...\n") |
| if ext == 'zip': |
| with zipfile.ZipFile(local_zip_path, "r") as zip_ref: |
| zip_ref.extractall(extract_to) |
| else: |
| with rarfile.RarFile(local_zip_path, "r") as zip_ref: |
| zip_ref.extractall(extract_to) |
|
|
| sys.stdout.write(">> Finished unzipping.\n") |
| os.remove(local_zip_path) |
|
|
| self.clear_statistics() |
|
|
| def compute_statistics(self, files): |
| """Use welford's method to compute mean and variance of the given |
| dataset. |
| |
| See https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online_algorithm.""" |
|
|
| assert len(files) > 1 |
|
|
| n = 0 |
| mean = np.zeros(3) |
| M2 = np.zeros(3) |
| for j, filename in enumerate(files): |
| |
| im = np.reshape(mpimg.imread(filename) * 255, [-1, 3]) |
| for i in range(np.shape(im)[1]): |
| n = n + 1 |
| delta = im[i] - mean |
| mean += delta / n |
| M2 += delta * (im[i] - mean) |
| sys.stdout.write('\r>> Processed %.1f%%' % ( |
| float(j) / float(len(files)) * 100.0)) |
| sys.stdout.flush() |
| var = M2 / (n - 1) |
| stddev = np.sqrt(var) |
| return np.float32(mean), np.float32(stddev) |
|
|
