def _save_patches(self,
patch_size=None,
stride=None,
force=False,
chunksize=10000):
if patch_size is None:
patch_save_dir = os.path.join(
self.data_dir, "_".join(self.all_vars),
'full_image-%i_%i' % (self.lr_km, self.hr_km))
else:
patches = True
patch_save_dir = os.path.join(
self.data_dir, "_".join(self.all_vars),
'patch_%i-%i_%i' % (patch_size, self.lr_km, self.hr_km))
if not os.path.exists(patch_save_dir):
os.makedirs(patch_save_dir)
files = []
for y in self.years:
print(y, 'Memory before inputs',
process.memory_info().rss / 10.**9)
patch_save_file = os.path.join(patch_save_dir,
str(y) + '_%i.tfrecords')
curr_files = glob.glob(os.path.join(patch_save_dir, str(y) + '*'))
if (len(curr_files) > 0) and (not force):
files += curr_files
continue
new_files = []
X, elevs, M, Y, lats, lons, t = self.get_patches(
y, patch_size, stride)
print "Shape of input", X.shape, "Shape of label", Y.shape
print "Writing files"
c = 0
fcount = 1
while c < len(X):
s = c + chunksize
new_files.append(patch_save_file % fcount)
convert_to_tf(X[c:s], elevs[c:s], M[c:s], Y[c:s], lats[c:s],
lons[c:s], t[c:s], new_files[-1])
c = s
fcount += 1
files += new_files
del X, elevs, Y, t, lats, lons
return files
def make_tf_test(self, save_file, scale1=1., scale2=0.5):
X, elev, Y, lats, lons, times = self.make_test(scale1, scale2)
order = range(Y.shape[0])
np.random.shuffle(order)
convert_to_tf(X[order], elev[order], Y[order], lats, lons,
times[order], save_file)
def make_patches(self,
save_file=None,
size=50,
stride=30,
lr_km=8,
hr_km=4,
shuffle=False):
if size is None:
return self.make_test(lr_km=lr_km, hr_km=hr_km)
factor = 1. * hr_km / lr_km
assert (size * factor) == int(size * factor)
assert (stride * factor) == int(stride * factor)
mask, da1, da2, elev = self.resolve_data(lr_km=lr_km, hr_km=hr_km)
obs_lats = da2.lat.values
obs_lons = da2.lon.values
X = da1.values
Y = da2.values
M = mask.values
M[np.isnan(M)] = 0
assert (np.max(np.abs(X)) < 1e20)
# keep elevation flexible by returning it seperately
elev = elev.values[:Y.shape[1], :Y.shape[2], np.newaxis]
X = np.expand_dims(X, 3)
labels, inputs, elevs, masks = [], [], [], []
lats, lons, times = [], [], []
timevals = da1.time.values
mask_vars = mask.notnull().values
land_locs = {}
for j, t in enumerate(timevals):
y = 0
x = 0
lasty = False
doubles = 0
while not lasty:
if (y + size) > Y.shape[1]:
y = Y.shape[1] - size
lasty = True
x = 0
lastx = False
while not lastx:
if (x + size) > Y.shape[2]:
x = Y.shape[2] - size
lastx = True
x_lr = int(x * factor)
y_lr = int(y * factor)
s_lr = int(size * factor)
x_sub = X[j, np.newaxis, y_lr:y_lr + s_lr,
x_lr:x_lr + s_lr]
# are we over the ocean?
#land_ratio = mask.notnull().values[y:y+size, x:x+size].mean()
if (x, y) not in land_locs.keys():
land_locs[(x, y)] = mask_vars[0, y:y + size,
x:x + size].mean() > 0.25
if land_locs[(x, y)]:
y_sub = Y[j, np.newaxis, y:y + size, x:x + size,
np.newaxis]
M_sub = M[j, np.newaxis, y:y + size, x:x + size,
np.newaxis]
elev_sub = elev[np.newaxis, y:y + size, x:x + size, :]
inputs += [x_sub]
labels += [y_sub]
elevs += [elev_sub]
masks += [M_sub]
lats += [obs_lats[np.newaxis, y:y + size]]
lons += [obs_lons[np.newaxis, x:x + size]]
times += [t]
x += stride
y += stride
order = range(len(inputs))
if shuffle:
np.random.shuffle(order)
self.inputs = np.concatenate(inputs, axis=0)[order]
self.labels = np.concatenate(labels, axis=0)[order]
masks = np.concatenate(masks, axis=0)[order]
elevs = np.concatenate(elevs, axis=0)[order]
lats = np.vstack(lats)[order]
lons = np.vstack(lons)[order]
if save_file is not None:
convert_to_tf(self.inputs, elevs, masks, self.labels, lats, lons,
np.array(times)[order], save_file)
return dict(inputs=self.inputs,
elevs=elevs,
labels=self.labels,
lats=lats,
lons=lons,
times=np.array(times)[order],
masks=masks)
def make_patches(self,
save_file=None,
size=50,
stride=30,
scale1=1.,
scale2=0.5):
assert (size * scale2) == int(size * scale2)
assert (stride * scale2) == int(stride * scale2)
mask, da1, da2, elev = self.resolve_data(scale1, scale2)
obs_lats = da2.lat.values
obs_lons = da2.lon.values
X = da1.values
Y = da2.values
# keep elevation flexible by returning it seperately
elev = elev.values[:Y.shape[1], :Y.shape[2], np.newaxis]
#tmp = np.empty(shape=(X.shape[0], X.shape[1], X.shape[2], 1))
#tmp[:] = elev
X = np.expand_dims(X, 3)
#X = np.concatenate([X, tmp], axis=3)
labels, inputs, elevs = [], [], []
lats, lons, times = [], [], []
timevals = da1.time.values
for j, t in enumerate(timevals):
for y in np.arange(0, Y.shape[1], stride):
for x in np.arange(0, Y.shape[2], stride):
if ((y + size) > Y.shape[1]) or ((x + size) > Y.shape[2]):
continue
x_lr = int(x * scale2)
y_lr = int(y * scale2)
s_lr = int(size * scale2)
x_sub = X[j, np.newaxis, y_lr:y_lr + s_lr,
x_lr:x_lr + s_lr]
# are we over the ocean?
land_ratio = mask.notnull().values[y:y + size,
x:x + size].mean()
if land_ratio < 0.50:
continue
y_sub = Y[j, np.newaxis, y:y + size, x:x + size,
np.newaxis]
elev_sub = elev[np.newaxis, y:y + size, x:x + size, :]
inputs += [x_sub]
labels += [y_sub]
elevs += [elev_sub]
lats += [obs_lats[np.newaxis, y:y + size]]
lons += [obs_lons[np.newaxis, x:x + size]]
times += [t]
order = range(len(inputs))
np.random.shuffle(order)
self.inputs = np.concatenate(inputs, axis=0)[order]
self.labels = np.concatenate(labels, axis=0)[order]
elevs = np.concatenate(elevs, axis=0)[order]
self.lats = np.vstack(lats)[order]
self.lons = np.vstack(lons)[order]
print "Number of subimages", len(self.inputs)
if save_file is not None:
convert_to_tf(self.inputs, elevs, self.labels, self.lats,
self.lons,
np.array(times)[order], save_file)