def __getitem__(self, index):
# load image
img = Image.open(self.image_filenames[index]).convert('RGB')
img = np.asarray(img)
img = util.modcrop(img, self.scale_factor)
# original HR image size
hr_img_w = img.shape[0]
hr_img_h = img.shape[1]
# determine lr_img LR image size
lr_img_w = hr_img_w // self.scale_factor
lr_img_h = hr_img_h // self.scale_factor
# only Y-channel is super-resolved
if self.is_gray:
img = np.asarray(img)
img = color.rgb2ycbcr(img) / 255
channel = len(img.shape)
img, _, _ = np.split(img, indices_or_sections=channel, axis=-1)
# img = img.convert('YCbCr')
# precision degrade from float64 to float32
img = Image.fromarray(img.squeeze())
# hr_img HR image
tensor_transform = tfs.ToTensor()
hr_img = tensor_transform(img)
# lr_img LR image
img = np.asarray(img)
if not self.bic_inp:
lr_img = imresize(img, output_shape=(lr_img_w, lr_img_h))
else:
lr_img = imresize(imresize(img, output_shape=(lr_img_w, lr_img_h)),
output_shape=(hr_img_w, hr_img_h))
bc_img = imresize(lr_img, output_shape=(hr_img_w, hr_img_h))
bc_img = Image.fromarray(bc_img.squeeze())
lr_img = Image.fromarray(lr_img.squeeze())
bc_img = tensor_transform(bc_img)
lr_img = tensor_transform(lr_img)
return lr_img, hr_img, bc_img
scaleFactor = 4.0
L = 1
lambda1 = 1e-4
net_idx = [
24, 23, 21, 19, 18, 16, 15, 14, 13, 12, 11, 10, 9, 8, 8, 7, 7, 6, 6, 5,
5, 4, 4, 4, 4, 3, 3, 3, 3, 2
]
maxIter = len(net_idx)
inIter = 5
C = 0.08
alpha = 1e-3
mu = 1e-3
K.clear_session()
HR = img_as_float(imread("./baby_GT.bmp")).astype(np.float32)
HR = modcrop(HR, int(scaleFactor))
if len(HR.shape) == 2:
HR = np.tile(HR[..., np.newaxis], (1, 1, 3))
HR_ycc = rgb2ycbcr(HR) / 255.0
label = HR_ycc[:, :, 0]
LR = imresize(label, 1.0 / scaleFactor, 'bicubic', 'F')
HR_bic = imresize(LR, scaleFactor, 'bicubic', 'F')
HR_ycc1 = copy.deepcopy(HR_ycc)
HR_ycc1[:, :, 0] = HR_bic
imsave('./lr.png', np.clip(ycbcr2rgb(HR_ycc1 * 255), 0.0, 1.0))
xlk = copy.deepcopy(HR_bic)
opt = parser.parse_args()
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
hr_list = [
join(opt.hr_set, x) for x in sorted(listdir(opt.hr_set))
if is_image_file(x)
] #glob.glob(opt.hr_set+"/*.bmp")
models = [load_model(path) for path in paths]
for id, image in enumerate(hr_list):
im_gt = io.imread(image)
img = color.rgb2ycbcr(im_gt) / 255
img = util.modcrop(img, opt.scale)
(rows, cols, channel) = img.shape
im_gt_y, im_bic_cb, im_bic_cr = np.split(img,
indices_or_sections=channel,
axis=2)
im_gt_y = im_gt_y.squeeze()
im_bic_cb = im_bic_cb.squeeze()
im_bic_cr = im_bic_cr.squeeze()
im_bic_cb = Image.fromarray(im_bic_cb.squeeze())
im_bic_cb = im_bic_cb.resize(
(int(cols // opt.scale), int(rows // opt.scale)),
resample=Image.BICUBIC)
im_bic_cb = im_bic_cb.resize((cols, rows), resample=Image.BICUBIC)
im_bic_cb = np.asarray(im_bic_cb, dtype=np.float64)
def sftgan(load_name="",
save_name='fin_rlt.png',
mode='rgb',
override_input=False):
path = load_name
test_img_folder_name = "TMP1"
# options
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
device = torch.device(
'cuda') # if you want to run on CPU, change 'cuda' -> 'cpu'
# device = torch.device('cpu')
# make dirs
test_img_folder = 'SFTGAN/data/' + test_img_folder_name # HR images
save_prob_path = 'SFTGAN/data/' + test_img_folder_name + '_segprob' # probability maps
save_byteimg_path = 'SFTGAN/data/' + test_img_folder_name + '_byteimg' # segmentation annotations
save_colorimg_path = 'SFTGAN/data/' + test_img_folder_name + '_colorimg' # segmentaion color results
util.mkdirs([save_prob_path, save_byteimg_path, save_colorimg_path])
test_prob_path = 'SFTGAN/data/' + test_img_folder_name + '_segprob' # probability maps
save_result_path = 'SFTGAN/data/' + test_img_folder_name + '_result' # results
util.mkdirs([save_result_path])
# load model
seg_model = arch.OutdoorSceneSeg()
seg_model_path = 'SFTGAN/pretrained_models/segmentation_OST_bic.pth'
seg_model.load_state_dict(torch.load(seg_model_path), strict=True)
seg_model.eval()
seg_model = seg_model.to(device)
# look_up table, RGB, for coloring the segmentation results
lookup_table = torch.from_numpy(
np.array([
[153, 153, 153], # 0, background
[0, 255, 255], # 1, sky
[109, 158, 235], # 2, water
[183, 225, 205], # 3, grass
[153, 0, 255], # 4, mountain
[17, 85, 204], # 5, building
[106, 168, 79], # 6, plant
[224, 102, 102], # 7, animal
[255, 255, 255], # 8/255, void
])).float()
lookup_table /= 255
print('Testing segmentation probability maps ...')
"""
for idx, path in enumerate(glob.glob(test_img_folder + '/*')):
imgname = os.path.basename(path)
basename = os.path.splitext(imgname)[0]
if "txt" in path:
continue
"""
idx = 0
if True:
#print(idx + 1, basename, path)
print(idx + 1)
# read image
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
img = util.modcrop(img, 8)
print(
"debug ",
img.shape,
img.ndim,
)
if img.ndim == 2:
img = np.expand_dims(img, axis=2)
if mode == 'bw':
#print(img.shape) # w,h,3 <- 1
stacked_img = np.stack((img, ) * 3, axis=2) # bw -> rgb
stacked_img = stacked_img[:, :, :, 0]
#print(stacked_img.shape) # w,h,3 <- 1
img = stacked_img
#(424, 1024, 3)
#print("debug img", img.shape, )
if override_input:
print("overriding input ", img.shape, "as", path)
util.save_img(img, path)
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
# MATLAB imresize
# You can use the MATLAB to generate LR images first for faster imresize operation
img_LR = util.imresize(img / 255, 1 / 4, antialiasing=True)
img = util.imresize(img_LR, 4, antialiasing=True) * 255
img[0] -= 103.939
img[1] -= 116.779
img[2] -= 123.68
img = img.unsqueeze(0)
img = img.to(device)
with torch.no_grad():
output = seg_model(img).detach().float().cpu().squeeze()
# save segmentation probability maps
#torch.save(output, os.path.join(save_prob_path, basename + '_bic.pth')) # 8xHxW
SEG_OUT = output
"""
# save segmentation byte images (annotations)
_, argmax = torch.max(output, 0)
argmax = argmax.squeeze().byte()
cv2.imwrite('foo1.png', argmax.numpy())
# save segmentation colorful results
im_h, im_w = argmax.size()
color = torch.FloatTensor(3, im_h, im_w).fill_(0) # black
for i in range(8):
mask = torch.eq(argmax, i)
color.select(0, 0).masked_fill_(mask, lookup_table[i][0]) # R
color.select(0, 1).masked_fill_(mask, lookup_table[i][1]) # G
color.select(0, 2).masked_fill_(mask, lookup_table[i][2]) # B
# void
mask = torch.eq(argmax, 255)
color.select(0, 0).masked_fill_(mask, lookup_table[8][0]) # R
color.select(0, 1).masked_fill_(mask, lookup_table[8][1]) # G
color.select(0, 2).masked_fill_(mask, lookup_table[8][2]) # B
torchvision.utils.save_image(
color, 'foo2.png', padding=0, normalize=False)
"""
del seg_model
'''
Codes for testing SFTGAN
'''
# options
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
sres_model_path = 'SFTGAN/pretrained_models/SFTGAN_torch.pth' # torch version
# sres_model_path = 'SFTGAN/pretrained_models/SFTGAN_noBN_OST_bg.pth' # pytorch version
device = torch.device(
'cuda') # if you want to run on CPU, change 'cuda' -> 'cpu'
# device = torch.device('cpu')
if 'torch' in sres_model_path: # torch version
model = arch.SFT_Net_torch()
else: # pytorch version
model = arch.SFT_Net()
model.load_state_dict(torch.load(sres_model_path), strict=True)
model.eval()
model = model.to(device)
print('Testing SFTGAN ...')
"""
for idx, path in enumerate(glob.glob(test_img_folder + '/*')):
imgname = os.path.basename(path)
basename = os.path.splitext(imgname)[0]
if "txt" in path:
continue
"""
if True:
path
#print(idx + 1, basename)
print(idx + 1)
# read image
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
img = util.modcrop(img, 8)
img = img * 1.0 / 255
if img.ndim == 2:
img = np.expand_dims(img, axis=2)
if mode == 'bw':
#print(img.shape) # w,h,3 <- 1
stacked_img = np.stack((img, ) * 3, axis=2) # bw -> rgb
stacked_img = stacked_img[:, :, :, 0]
#print(stacked_img.shape) # w,h,3 <- 1
img = stacked_img
#(424, 1024, 3)
#print("debug img", img.shape, )
img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]],
(2, 0, 1))).float()
# MATLAB imresize
# You can use the MATLAB to generate LR images first for faster imresize operation
img_LR = util.imresize(img, 1 / 4, antialiasing=True)
img_LR = img_LR.unsqueeze(0)
img_LR = img_LR.to(device)
# read segmentation probability maps
#seg = torch.load(os.path.join(test_prob_path, basename + '_bic.pth'))
seg = SEG_OUT
seg = seg.unsqueeze(0)
# change probability
# seg.fill_(0)
# seg[:,5].fill_(1)
seg = seg.to(device)
with torch.no_grad():
output = model((img_LR, seg)).data.float().cpu().squeeze()
output = util.tensor2img(output)
util.save_img(output, save_name)
model = arch.SFT_Net_torch()
else: # pytorch version
model = arch.SFT_Net()
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
print('Testing SFTGAN ...')
for idx, path in enumerate(glob.glob(test_img_folder + '/*')):
imgname = os.path.basename(path)
basename = os.path.splitext(imgname)[0]
print(idx + 1, basename)
# read image
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
img = util.modcrop(img, 8)
img = img * 1.0 / 255
if img.ndim == 2:
img = np.expand_dims(img, axis=2)
img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]],
(2, 0, 1))).float()
# MATLAB imresize
# You can use the MATLAB to generate LR images first for faster imresize operation
img_LR = util.imresize(img, 1 / 4, antialiasing=True)
img_LR = img_LR.unsqueeze(0)
img_LR = img_LR.to(device)
# read segmentation probability maps
seg = torch.load(os.path.join(test_prob_path, basename + '_bic.pth'))
seg = seg.unsqueeze(0)
# change probability
lr_bic_image_list = []
stride = 7
for f in datafiles:
img = cv2.imread(FLAGS.image_folder + f).astype(np.float32)
if img.ndim == 3:
img = img[:, :, ::-1] # BGR to RGB
img_ycbcr = util.rgb2ycbcr(img / 255.0) * 255.0
# img_y = img_ycbcr[:, :, 0]
# img_y = np.rint(np.clip(img_y, 0, 255))
# img_y = util.modcrop(img_y, FLAGS.scale)
# img_y_l = imresize.fast_imresize(img_y, 1 / float(FLAGS.scale))
# img_y_b = imresize.fast_imresize(img_y_l, FLAGS.scale)
# # img_y_b = np.rint(np.clip(img_y_b, 0, 255))
img_ycbcr = util.modcrop(img_ycbcr, FLAGS.scale)
img_y = img_ycbcr[:, :, 0]
img_y = np.rint(np.clip(img_y, 0, 255))
img_ycbcr_l = imresize.fast_imresize(img_ycbcr,
1 / float(FLAGS.scale))
img_ycbcr_b = imresize.fast_imresize(img_ycbcr_l, FLAGS.scale)
img_y_b = img_ycbcr_b[:, :, 0]
lr_bic_image_list.append(img_ycbcr_b)
else:
img_y = util.modcrop(img, FLAGS.scale)
img_y_l = imresize.fast_imresize(img_y, 1 / float(FLAGS.scale))
img_y_b = imresize.fast_imresize(img_y_l, FLAGS.scale)
lr_bic_image_list.append(img_y_b)
lr_image_list.append(img_y_b.astype(np.float32))
hr_image_list.append(img_y.astype(np.float32))