File size: 5,099 Bytes
872b0a0 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """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():
# Should be a list containing all subdirectories of the main data dir which
# belong to this dataset
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
#if self.fast_dir:
# print(">> Removing files from {}".format(self.fast_dir))
# for d in self.dirs:
# shutil.rmtree(os.path.join(self.fast_dir, d))
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):
#TODO ensure the pixel values are 0..255
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)
|