class EnhancementOut(object):
def __init__(self, network_out: prob_clf.NetworkOutput,
x_r: NormalizedTensor, x_l: NormalizedTensor):
self.network_out = network_out
self.x_r = x_r.to_sym()
res = self.x_r.t - x_l.to_sym().t
self.res_sym = SymbolTensor(res, L=511, centered=True)
self.res = self.res_sym.to_norm()
self._mean_img = None
def get_mean_img(self, loss: DiscretizedMixLogisticLoss):
if self._mean_img is None:
self._mean_img = extract_mean_image_corrected(
self.res, self.network_out, loss)
return self._mean_img
def unpack(self, img_batch, fixed_first=None) -> InputTensors:
raw = img_batch['raw'].to(pe.DEVICE,
non_blocking=True) # uint8 or int16
compressed = img_batch['compressed'].to(
pe.DEVICE, non_blocking=True) # uint8 or int16
bpps = img_batch['bpp'].to(pe.DEVICE) # 1d tensor of floats
if fixed_first is not None:
raw[0, ...] = fixed_first['raw']
compressed[0, ...] = fixed_first['compressed']
bpps[0] = fixed_first['bpp']
num_subpixels_before_pad = np.prod(raw.shape)
if self.padding_fac:
raw = self.pad(raw)
compressed = self.pad(compressed)
s_c = SymbolTensor(compressed.long(), L=256)
x_c = s_c.to_norm()
s_r = SymbolTensor(raw.long(), L=256)
x_r = s_r.to_norm()
return InputTensors((x_r, x_c), bpps, num_subpixels_before_pad)
def pad_pack(self, raw, compressed, bpps) -> InputTensors:
""" Pad iimages and pack into a InputTensors instance
:param raw: Batch of raw input images
:param compressed: Output of compressing images in `raw` with BPG.
:param bpps: The bitrates of the images.
:return: InputTensors
"""
assert raw.shape == compressed.shape
num_subpixels_before_pad = np.prod(raw.shape)
if self.padding_fac:
raw = self.pad(raw)
compressed = self.pad(compressed)
assert len(raw.shape) == 4
assert len(bpps.shape) == 1, (bpps.shape, raw.shape)
s_c = SymbolTensor(compressed.long(), L=256)
x_c = s_c.to_norm()
s_r = SymbolTensor(raw.long(), L=256)
x_r = s_r.to_norm()
return InputTensors((x_r, x_c), bpps, num_subpixels_before_pad)
def unroll_unpack(self, img):
"""
Main function. Given an image `img`. Unrolls _unpack over all needed Qs:
- if the dataset is not meta:
just return _unpack(image)
- otherwise:
if run_clf:
set self.clf_q from classifier
if not CLF_ONLY:
do the parabola search.
Always yields a set of images, and the q used to optain it.
"""
if not self.is_meta:
yield self._unpack(img), None
else:
# get a first Q guess
if self.run_clf:
assert self.clf
# (x_r, _), _, _ = self._unpack(img[next(iter(img))])
raw, _, _ = self._unpack(img[next(iter(img))])
s_r = SymbolTensor(raw.long(), L=256)
x_r = s_r.to_norm()
torch.cuda.empty_cache()
# crop_qs = [self.clf.get_q(x_r_crop) for x_r_crop in auto_crop.iter_crops(x_r.get())]
# print('***\n',crop_qs,'\n')
self.clf_q = self.clf.get_q(x_r.get())
torch.cuda.empty_cache()
elif self.qstrategy == QStrategy.FIXED:
# note: we use clf_q also for the fixed Q, confusing naming scheme, I know!!!
self.clf_q = 14 # optimal on training set
# do optimum find
if self.qstrategy != QStrategy.CLF_ONLY:
while self.next_q and self.next_q in img:
# this is a dict with raw, compressed, bpps now
yield self._unpack(img[self.next_q]), self.next_q
# make sure the guess was also evaluated
if self.clf_q and self.clf_q not in self.losses:
if self.clf_q not in img:
raise ValueError('**** CLF returned invalid q', self.clf_q)
else:
# get this one as well
yield self._unpack(img[self.clf_q]), self.clf_q
def unpack(self, img_batch, fixed_first=None):
raw = img_batch['raw'].to(pe.DEVICE,
non_blocking=True) # uint8 or int16
q = img_batch['q'].to(pe.DEVICE).view(-1) # 1d tensor of floats
if fixed_first is not None:
raw[0, ...] = fixed_first['raw']
q[0, ...] = fixed_first['q']
q = q - self.config_clf.first_class
if self.padding_fac:
raw = self.pad(raw)
raw = SymbolTensor(raw.long(), L=256)
return raw.to_norm(), q
def unpack_batch_pad(self, img_or_imgbatch):
raw = img_or_imgbatch['raw'].to(pe.DEVICE,
non_blocking=True) # uint8 or int16
q = img_or_imgbatch['q'].to(pe.DEVICE).view(-1) # 1d tensor of floats
if len(raw.shape) == 3:
raw.unsqueeze_(0)
if self.padding_fac:
raw = self.pad(raw)
q = q - self.config_clf.first_class
assert len(raw.shape) == 4
raw = SymbolTensor(raw.long(), L=256)
return raw.to_norm(), q