def _generate_crop_grid(self, frames, size, shuffle=False):
"""generate randomly cropped box of `frames`
Args:
frames: a list of tuple, commonly returned from `_process_at_file`.
size: an int scalar to specify number of generated crops.
shuffle: a boolean, whether to shuffle the outputs.
Return:
list of tuple: containing (HR, LR, box, name) respectively, where HR and LR
are reference frames, box is a list of 4 int of crop coordinates.
"""
patch_size = Utility.to_list(self.patch_size, 2)
patch_size = Utility.shrink_mod_scale(patch_size, self.scale)
if size < 0:
index = np.arange(len(frames)).tolist()
else:
if self.crop == 'random':
index = np.random.randint(len(frames), size=size).tolist()
else:
index = np.arange(size).tolist()
grids = []
for i in range(len(frames)):
hr, lr, name = frames[i]
_w, _h = hr[0].width, hr[0].height
if self.crop == 'not' or self.crop is None:
_pw, _ph = _w, _h
else:
_pw, _ph = patch_size
amount = index.count(i)
if self.crop == 'random':
x = np.random.randint(0, _w - _pw + 1, size=amount)
y = np.random.randint(0, _h - _ph + 1, size=amount)
elif self.crop == 'center':
x = np.array([(_w - _pw) // 2] * amount)
y = np.array([(_h - _ph) // 2] * amount)
else:
x = np.zeros([amount])
y = np.zeros([amount])
x -= x % self.scale[0]
y -= y % self.scale[1]
grids += [(ImageProcess.img_to_array(hr),
ImageProcess.img_to_array(lr),
[_x, _y, _x + _pw, _y + _ph], name)
for _x, _y in zip(x, y)]
if shuffle:
np.random.shuffle(grids)
return grids
def test_resize_upsample():
Im = Image.open(URL)
for X in [Im, Im.convert('L')]:
w = X.width
h = X.height
for ss in [2, 3, 4, 5, 6]:
GT = X.resize([w * ss, h * ss], Image.BICUBIC)
gt = np.asarray(GT, dtype='float32') / 255
x = tf.constant(np.asarray(X), dtype='float32') / 255
y = U.upsample(x, ss).numpy().clip(0, 1)
assert np.all(np.abs(y[16:-16, 16:-16] - gt[16:-16, 16:-16]) < 1.0e-2), \
f"Scale: {ss}. Mode: {X.mode}"
def test_resize_2x2():
scale = 2
img = np.array([[1, 2], [3, 4]], np.uint8)
imgp = Image.fromarray(img, 'L')
img = np.reshape(img, [1, 2, 2, 1])
y = U.upsample(img.astype(np.float32), scale)
y = y.numpy()[0, ..., 0].astype('uint8')
imgp = imgp.resize([2 * scale, 2 * scale], Image.BICUBIC)
yp = np.array(imgp)
yf = tf.image.resize_bicubic(img, [2 * scale, 2 * scale])
yf = yf.numpy()[0, ..., 0].astype('uint8')
diff = yp - y
def test_resize_img():
scale = 2
img = Image.open(URL)
w = img.width
h = img.height
imgp = img.resize([w * scale, h * scale], Image.BICUBIC)
yp = np.array(imgp)
img = np.resize(img, [1, h, w, 3])
y = U.upsample(img.astype(np.float32), scale)
y = y.numpy()[0].astype('uint8')
yf = tf.image.resize_bicubic(img, [h * scale, w * scale])
yf = yf.numpy()[0].astype('uint8')
diff = yp - y
def _parse_config(self, config, **kwargs):
assert isinstance(config, Config)
config.update(kwargs)
_needed_args = ('batch', 'depth', 'patch_size', 'scale',
'steps_per_epoch', 'convert_to')
for _arg in _needed_args:
if _arg not in config:
raise ValueError(_arg +
' is required in config, but not found.')
self.depth = config.depth
self.patch_size = config.patch_size
self.scale = Utility.to_list(config.scale, 2)
self.patches_per_epoch = config.steps_per_epoch * config.batch
self.batch = config.batch
self.crop = config.crop
def test_resize_downsample():
Im = Image.open(URL)
for X in [Im, Im.convert('L')]:
w = X.width
h = X.height
for ss in [2, 4, 6, 8]:
w_ = w - w % ss
h_ = h - h % ss
X = X.crop([0, 0, w_, h_])
GT = X.resize([w_ // ss, h_ // ss], Image.BICUBIC)
gt = np.asarray(GT, dtype='float32') / 255
x = tf.constant(np.asarray(X), dtype='float32') / 255
y = U.downsample(x, ss).numpy().clip(0, 1)
assert np.all(np.abs(y[8:-8, 8:-8] - gt[8:-8, 8:-8]) < 1.0e-2), \
f"Scale: {ss}. Mode: {X.mode}"
def _prefetch(self, memory_usage=None, shard=1, index=0):
"""Prefetch `size` files and load into memory. Specify `shard` will divide
loading files into `shard` shards in order to execute in parallel.
NOTE: parallelism is implemented via `QuickLoader`
Args:
memory_usage: desired virtual memory to use, could be int (in bytes) or
a readable string (i.e. '3GB', '1TB'). Default to use all available
memories.
shard: an int scalar to specify the number of shards operating in parallel.
index: an int scalar, representing shard index
"""
if self.all_loaded:
interval = len(self.frames) // shard
return self.frames[index * interval:(index + 1) * interval], True
# check memory usage
if isinstance(memory_usage, str):
memory_usage = Utility.str_to_bytes(memory_usage)
if not memory_usage:
memory_usage = virtual_memory().total
memory_usage = np.min([np.uint64(memory_usage), virtual_memory().free])
capacity = self.size
frames = []
if capacity <= memory_usage and shard == 1 or capacity <= memory_usage // 2:
# load all clips
self.all_loaded = True
interval = len(self.file_objects) // shard
for file in self.file_objects[index * interval:(index + 1) *
interval]:
frames += self._process_at_file(file, file.frames)
else:
prop = memory_usage / capacity / shard / 2
size = int(np.round(len(self) * prop))
for file, amount in self._random_select(size).items():
frames += self._process_at_file(file, amount)
self.frames = frames
return frames, self.all_loaded
def dummy_test_str_to_bytes():
for t, a in zip(TEST_STR, ANS):
ans = U.str_to_bytes(t)
print(t, ans)
assert ans == a
def build_loader(self, crop=True, **kwargs):
"""Build image1(s) pair loader, make self iterable
Args:
crop: if True, crop the images into patches
kwargs: you can override attribute in the dataset
"""
_crop_args = ['scale', 'patch_size', 'strides', 'depth']
for _arg in _crop_args:
if _arg in kwargs and kwargs[_arg]:
self.__setattr__(_arg, kwargs[_arg])
self.scale = Utility.to_list(self.scale, 2)
self.patch_size = Utility.to_list(self.patch_size, 2)
self.strides = Utility.to_list(self.strides, 2)
self.patch_size = Utility.shrink_mod_scale(
self.patch_size, self.scale) if crop else None
self.strides = Utility.shrink_mod_scale(self.strides,
self.scale) if crop else None
for vf in self.dataset:
tf.logging.debug('loading ' + vf.name)
depth = self.depth
# read all frames if depth is set to -1
if depth == -1: depth = vf.frames
for _ in range(vf.frames // depth):
frames_hr = [
ImageProcess.shrink_to_multiple_scale(img, self.scale)
if self.modcrop else img for img in vf.read_frame(depth)
]
frames_lr = [
ImageProcess.imresize(img,
np.ones(2) / self.scale)
for img in frames_hr
]
self.frames.append((frames_hr, frames_lr, vf.name))
vf.reopen()
tf.logging.debug(
'all files load finished, generating cropping meshes...')
self.random = self.random and crop
if self.random:
rand_index = np.random.randint(len(self.frames),
size=self.max_patches)
for i in rand_index:
hr, lr, name = self.frames[i]
_w, _h = hr[0].width, hr[0].height
_pw, _ph = self.patch_size or [_w, _h]
x = np.random.randint(0, _w - _pw + 1)
x -= x % self.scale[0]
y = np.random.randint(0, _h - _ph + 1)
y -= y % self.scale[1]
self.grid.append((hr, lr, x, y, name))
else:
for hr, lr, name in self.frames:
_w, _h = hr[0].width, hr[0].height
_sw, _sh = self.strides or [_w, _h]
_pw, _ph = self.patch_size or [_w, _h]
x, y = np.mgrid[0:_w - _pw - (_w - _pw) % _sw + _sw:_sw,
0:_h - _ph - (_h - _ph) % _sh + _sh:_sh]
self.grid += [(hr, lr, _x, _y, name)
for _x, _y in zip(x.flatten(), y.flatten())]
tf.logging.info('data loader is ready!')
self.batch_iterator = self._build_iter()
self.built = True