def __getitem__(self, idx):
# load GT image
pathGT = os.path.join(self.imgGTRoot, self.imgNames[idx])
imgGT = util.imread_uint8(pathGT, mode=self.colorMode)
# ========== High quality patch ==========
imgGT = util.random_crop(imgGT, self.size)
imgGT = util.augment_img(imgGT)
# ========== Low quality patch ==========
imgGT = Image.fromarray(imgGT)
assert self.quality[0] <= self.quality[
1], f'Illegal quality range {self.quality}'
quality = random.randint(self.quality[0], self.quality[1])
with BytesIO() as buffer:
imgGT.save(buffer, format='JPEG', quality=quality)
imgIn = Image.open(buffer).convert()
imgGT = np.array(imgGT, copy=False)
imgIn = np.array(imgIn, copy=False)
imgIn = util.uint2tensor(imgIn)
imgGT = util.uint2tensor(imgGT)
return imgIn, imgGT
def __getitem__(self, idx):
# load GT image
pathGT = os.path.join(self.imgGTRoot, self.imgNames[idx])
imgGT = util.imread_uint8(pathGT, mode=self.colorMode)
# ========== High quality patch ==========
imgGT = util.random_crop(imgGT, self.size)
imgGT = util.augment_img(imgGT)
# ========== Low quality patch ==========
imgGT = Image.fromarray(imgGT)
assert self.quality[0] <= self.quality[1], f'Illegal quality range {self.quality}'
quality = random.randint(self.quality[0], self.quality[1])
with BytesIO() as buffer:
imgGT.save(buffer, format='JPEG', quality=quality)
QMs = getQM(buffer.getvalue())
imgIn = Image.open(buffer).convert()
imgGT = np.array(imgGT, copy=False)
imgIn = np.array(imgIn, copy=False)
# ========== Quantization map ==========
if self.colorMode == 'L':
QM = np.zeros([8,8,1], dtype=np.uint8)
QM[:,:,0] = QMs[0]
else:
QM = np.zeros([8,8,2], dtype=np.uint8)
QM[:,:,0] = QMs[0]
QM[:,:,1] = QMs[1]
assert self.size % 8 == 0, f'The arg \"size\" should be times of 8, but {self.size}.'
cnt = self.size // 8
QMimg = np.tile(QM, (cnt,cnt,1))
if imgIn.ndim == 2:
imgIn = imgIn[:, :, np.newaxis]
if imgGT.ndim == 2:
imgGT = imgGT[:, :, np.newaxis]
imgGlobal = skimage.transform.resize(imgIn, (112,112), order=3)
imgIn = np.concatenate((imgIn, QMimg), axis=-1)
imgIn = util.uint2tensor(imgIn)
imgGT = util.uint2tensor(imgGT)
imgGlobal = util.uint2tensor(imgGlobal)
return [imgIn, imgGlobal], imgGT
def __getitem__(self, idx):
# load GT image
pathGT = os.path.join(self.imgGTRoot, self.imgNames[idx])
imgGT = util.imread_uint8(pathGT, mode=self.colorMode)
# ========== Low quality image ==========
imgGT = Image.fromarray(imgGT)
with BytesIO() as buffer:
imgGT.save(buffer, format='JPEG', quality=self.quality)
imgIn = Image.open(buffer).convert()
imgGT = np.array(imgGT, copy=False)
imgIn = np.array(imgIn, copy=False)
imgIn = util.uint2tensor(imgIn)
imgGT = util.uint2tensor(imgGT)
return imgIn, imgGT
def __getitem__(self, idx):
# load GT image
pathGT = os.path.join(self.imgGTRoot, self.imgNames[idx])
imgGT = util.imread_uint8(pathGT, mode=self.colorMode)
# pad the image to multiple of 8
ori_h, ori_w = imgGT.shape[0:2]
imgGT = self.padding8(imgGT)
# ========== Low quality image ==========
imgGT = Image.fromarray(imgGT)
with BytesIO() as buffer:
imgGT.save(buffer, format='JPEG', quality=self.quality)
QMs = getQM(buffer.getvalue())
imgIn = Image.open(buffer).convert()
imgGT = np.array(imgGT, copy=False)
imgIn = np.array(imgIn, copy=False)
LQimg = imgIn.copy()
# ========== Quantization map ==========
if self.colorMode == 'L':
QM = np.zeros([8,8,1], dtype=np.uint8)
QM[:,:,0] = QMs[0]
else:
QM = np.zeros([8,8,2], dtype=np.uint8)
QM[:,:,0] = QMs[0]
QM[:,:,1] = QMs[1]
cnt_h, cnt_w = np.array(imgIn.shape[0:2]) // 8
QMimg = np.tile(QM, (cnt_h, cnt_w, 1))
if imgIn.ndim == 2:
imgIn = imgIn[:, :, np.newaxis]
if imgGT.ndim == 2:
imgGT = imgGT[:, :, np.newaxis]
imgGlobal = skimage.transform.resize(imgIn, (112,112), order=3)
imgIn = np.concatenate((imgIn, QMimg), axis=-1)
imgIn = util.uint2tensor(imgIn)
imgGT = util.uint2tensor(imgGT)
imgGlobal = util.uint2tensor(imgGlobal)
return [imgIn, imgGlobal], imgGT, LQimg, (ori_h, ori_w)
dataloader = trainOpt['dataloader'] if 'dataloader' in trainOpt else ''
valid_set = eval(dataloader).ValidData(imgValidRoot=imgGTRoot,
quality=opt.quality,
colorMode=opt.colorMode)
for i in range(len(valid_set)):
currStartTime = time.time()
print(f'\n[{i+1}/{len(valid_set)}]', valid_set.imgNames[i])
if opt.GT:
currImg, GTImg, LQImg, inSize = valid_set.__getitem__(i)
currImg = util.tensor3to4(currImg)
GTImg = util.tensor2uint(GTImg)
else:
#TODO: add QM
currImg = util.imread_uint8(fullPaths[i], mode=opt.colorMode)
currImg = util.uint2tensor(currImg)
currImg = util.tensor3to4(currImg)
start_time = time.time()
if cuda:
currImg = currImg.cuda()
HQImg = model(currImg)
HQImg = util.tensor2uint(HQImg)
elapsed_time = time.time() - start_time
gputime += elapsed_time
if opt.GT:
LQImg = LQImg[0:inSize[0], 0:inSize[1], ...]
HQImg = HQImg[0:inSize[0], 0:inSize[1], ...]