def __getitem__(self, index):
path = os.path.join(self.data_path, self.samples[index])
img = default_loader(path)
img = transforms.ToTensor()(img) # turn the image to a tensor
if self.isTrain == 'test':
gt_path = os.path.join(self.gt_path, self.samples[index])
mask = default_loader(gt_path)
mask = transforms.ToTensor()(mask) # turn the image to a tensor
return img, mask
return img
def test_contextual_attention(args):
import cv2
import os
# run on cpu
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
def float_to_uint8(img):
img = img * 255
return img.astype('uint8')
rate = 2
stride = 1
grid = rate * stride
b = default_loader(args.imageA)
w, h = b.size
b = b.resize((w // grid * grid // 2, h // grid * grid // 2),
Image.ANTIALIAS)
# b = b.resize((w//grid*grid, h//grid*grid), Image.ANTIALIAS)
print('Size of imageA: {}'.format(b.size))
f = default_loader(args.imageB)
w, h = f.size
f = f.resize((w // grid * grid, h // grid * grid), Image.ANTIALIAS)
print('Size of imageB: {}'.format(f.size))
f, b = transforms.ToTensor()(f), transforms.ToTensor()(b)
f, b = f.unsqueeze(0), b.unsqueeze(0)
if torch.cuda.is_available():
f, b = f.cuda(), b.cuda()
contextual_attention = ContextualAttention(ksize=3,
stride=stride,
rate=rate,
fuse=True)
if torch.cuda.is_available():
contextual_attention = contextual_attention.cuda()
yt, flow_t = contextual_attention(f, b)
vutils.save_image(yt, 'vutils' + args.imageOut, normalize=True)
vutils.save_image(flow_t, 'flow' + args.imageOut, normalize=True)
def __getitem__(self, index):
path = os.path.join(self.data_path, self.samples[index])
img = default_loader(path)
if self.random_crop:
imgw, imgh = img.size
if imgh < self.image_shape[0] or imgw < self.image_shape[1]:
img = transforms.Resize(min(self.image_shape))(img)
img = transforms.RandomCrop(self.image_shape)(img)
else:
img = transforms.Resize(self.image_shape)(img)
img = transforms.RandomCrop(self.image_shape)(img)
img = transforms.ToTensor()(img) # turn the image to a tensor
img = normalize(img)
if self.return_name:
return self.samples[index], img
else:
return img
def generate(img, img_mask_path, model_path):
with torch.no_grad(): # enter no grad context
if img_mask_path and is_image_file(img_mask_path):
# Test a single masked image with a given mask
x = Image.fromarray(img)
mask = default_loader(img_mask_path)
x = transforms.Resize(config['image_shape'][:-1])(x)
x = transforms.CenterCrop(config['image_shape'][:-1])(x)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(config['image_shape'][:-1])(mask)
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
elif img_mask_path:
raise TypeError("{} is not an image file.".format(img_mask_path))
else:
# Test a single ground-truth image with a random mask
#ground_truth = default_loader(img_path)
ground_truth = img
ground_truth = transforms.Resize(config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.CenterCrop(config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.ToTensor()(ground_truth)
ground_truth = normalize(ground_truth)
ground_truth = ground_truth.unsqueeze(dim=0)
bboxes = random_bbox(config, batch_size=ground_truth.size(0))
x, mask = mask_image(ground_truth, bboxes, config)
# Set checkpoint path
if not model_path:
checkpoint_path = os.path.join('checkpoints',
config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = model_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
last_model_name = get_model_list(checkpoint_path, "gen", iteration=0)
if cuda:
netG.load_state_dict(torch.load(last_model_name))
else:
netG.load_state_dict(torch.load(last_model_name, map_location='cpu'))
model_iteration = int(last_model_name[-11:-3])
print("Resume from {} at iteration {}".format(checkpoint_path, model_iteration))
if cuda:
netG = nn.parallel.DataParallel(netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result = from_torch_img_to_numpy(inpainted_result, 'output.png', padding=0, normalize=True)
return inpainted_result
def main():
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
print("Arguments: {}".format(args))
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
print("Configuration: {}".format(config))
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(args.image):
if args.mask and is_image_file(args.mask):
# Test a single masked image with a given mask
x = default_loader(args.image)
mask = default_loader(args.mask)
x = transforms.Resize(config['image_shape'][:-1])(x)
x = transforms.CenterCrop(config['image_shape'][:-1])(x)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(
config['image_shape'][:-1])(mask)
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
elif args.mask:
raise TypeError(
"{} is not an image file.".format(args.mask))
else:
# Test a single ground-truth image with a random mask
ground_truth = default_loader(args.image)
ground_truth = transforms.Resize(
config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.CenterCrop(
config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.ToTensor()(ground_truth)
ground_truth = normalize(ground_truth)
ground_truth = ground_truth.unsqueeze(dim=0)
bboxes = random_bbox(
config, batch_size=ground_truth.size(0))
x, mask = mask_image(ground_truth, bboxes, config)
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join('checkpoints',
config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
last_model_name = get_model_list(
checkpoint_path, "gen", iteration=args.iter)
netG.load_state_dict(torch.load(last_model_name))
model_iteration = int(last_model_name[-11:-3])
print("Resume from {} at iteration {}".format(
checkpoint_path, model_iteration))
if cuda:
netG = nn.parallel.DataParallel(
netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
vutils.save_image(inpainted_result, args.output,
padding=0, normalize=True)
print("Saved the inpainted result to {}".format(args.output))
if args.flow:
vutils.save_image(offset_flow, args.flow,
padding=0, normalize=True)
print("Saved offset flow to {}".format(args.flow))
else:
raise TypeError("{} is not an image file.".format)
# exit no grad context
except Exception as e: # for unexpected error logging
print("Error: {}".format(e))
raise e
def main():
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
chunker = ImageChunker(config['image_shape'][0], config['image_shape'][1],
args.overlap)
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(args.image):
print("Loading image...")
imgs, masks = [], []
img_ori = default_loader(args.image)
img_w, img_h = img_ori.size
# Load mask txt file
fname = args.image.replace('.jpg', '.txt')
bboxes, _ = load_bbox_txt(fname, img_w, img_h)
mask_ori = create_mask(bboxes, img_w, img_h)
chunked_images = chunker.dimension_preprocess(
np.array(deepcopy(img_ori)))
chunked_masks = chunker.dimension_preprocess(
np.array(deepcopy(mask_ori)))
for (x, msk) in zip(chunked_images, chunked_masks):
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(msk)[0].unsqueeze(dim=0)
# x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
imgs.append(x)
masks.append(mask)
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
last_model_name = get_model_list(checkpoint_path,
"gen",
iteration=args.iter)
netG.load_state_dict(torch.load(last_model_name))
model_iteration = int(last_model_name[-11:-3])
print("Resume from {} at iteration {}".format(
checkpoint_path, model_iteration))
pred_imgs = []
for (x, mask) in zip(imgs, masks):
if torch.max(mask) == 1:
if cuda:
netG = nn.parallel.DataParallel(
netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result = inpainted_result.squeeze(
dim=0).permute(1, 2, 0).cpu()
pred_imgs.append(inpainted_result.numpy())
else:
pred_imgs.append(
x.squeeze(dim=0).permute(1, 2, 0).numpy())
pred_imgs = np.asarray(pred_imgs, dtype=np.float32)
reconstructed_image = chunker.dimension_postprocess(
pred_imgs, np.array(img_ori))
# plt.imshow(reconstructed_image); plt.show()
reconstructed_image = torch.tensor(
reconstructed_image).permute(2, 0, 1).unsqueeze(dim=0)
vutils.save_image(reconstructed_image,
args.output,
padding=0,
normalize=True)
print("Saved the inpainted result to {}".format(args.output))
if args.flow:
vutils.save_image(offset_flow,
args.flow,
padding=0,
normalize=True)
print("Saved offset flow to {}".format(args.flow))
else:
raise TypeError("{} is not an image file.".format)
# exit no grad context
except Exception as e: # for unexpected error logging
print("Error: {}".format(e))
raise e
def generateInpaintedImage(args, netG, imagePath):
config = get_config(args.g_config)
occlusions = []
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
print("Arguments: {}".format(args))
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(imagePath):
if args.mask and is_image_file(args.mask):
# Test a multiple masked image with a given mask
x = default_loader(imagePath)
x = transforms.Resize([512, 1024])(x)
mask = default_loader(args.mask)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(
config['image_shape'][:-1])(mask)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
w, h = x.size
first = x.crop((0, 0, w // 3, h))
second = x.crop((w // 3, 0, ((w // 3) * 2) + 2, h))
third = x.crop(((w // 3) * 2, 0, w, h))
for y in [first, second, third]:
y = transforms.CenterCrop(
config['image_shape'][:-1])(y)
y = transforms.ToTensor()(y)
y = normalize(y)
y = y * (1. - mask)
occlusions.append(y)
elif args.mask:
raise TypeError("{} is not an image file.".format(
args.mask))
default_image = default_loader(imagePath)
di_w, di_h = default_image.size
for idx, occlusion in enumerate(occlusions):
if cuda:
occlusion = occlusion.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(occlusion, mask)
inpainted_result = x2 * mask + occlusion * (1. - mask)
inp_hw = config['image_shape'][1]
if idx == 0:
offset = ((di_w // 3 - inp_hw) // 2,
(di_h - inp_hw) // 2)
elif idx == 1:
offset = ((di_w - inp_hw) // 2, (di_h - inp_hw) // 2)
elif idx == 2:
offset = ((((di_w - inp_hw) // 2) + (di_w // 3)),
(di_h - inp_hw) // 2)
grid = vutils.make_grid(inpainted_result, normalize=True)
# Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute(
1, 2, 0).to('cpu', torch.uint8).numpy()
im = Image.fromarray(ndarr)
im = transforms.CenterCrop(config['mask_shape'])(im)
im = transforms.Resize(config['image_shape'][:-1])(im)
default_image.paste(im, offset)
return default_image
else:
raise TypeError("{} is not an image file.".format)
# exit no grad context
except Exception as e: # for unexpected error logging
print("Error: {}".format(e))
raise e
def main():
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
print("Arguments: {}".format(args))
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
print("Configuration: {}".format(config))
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
file = dataset_files(args.test_root, "*.jpg")
mask_file = dataset_files(args.mask_dir, "*.png")
for j in range(len(mask_file)):
for i in range(len(file)):
if is_image_file(file[i]):
if mask_file and is_image_file(mask_file[j]):
# Test a single masked image with a given mask
x = default_loader(file[i])
mask = default_loader(mask_file[j])
# x = cv2.cvtColor(cv2.imread(file[i]), cv2.COLOR_BGR2RGB)
# mask = cv2.cvtColor(cv2.imread(mask_file[j]), cv2.COLOR_BGR2RGB)
# x = cv2.resize(x, (config['image_shape'][0], config['image_shape'][1]))
# mask = cv2.resize(mask, (config['image_shape'][0], config['image_shape'][1]))
x = transforms.Resize(
config['image_shape'][:-1])(x)
x = transforms.CenterCrop(
config['image_shape'][:-1])(x)
# mask = transforms.Resize(config['image_shape'][:-1])(mask)
# mask = transforms.CenterCrop(config['image_shape'][:-1])(mask)
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(mask)[0].unsqueeze(
dim=0)
x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
# x_raw = x
mask = mask.unsqueeze(dim=0)
elif mask_file[j]:
raise TypeError("{} is not an image file.".format(
mask_file[j]))
else:
# Test a single ground-truth image with a random mask
ground_truth = default_loader(file[i])
ground_truth = transforms.Resize(
config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.CenterCrop(
config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.ToTensor()(ground_truth)
ground_truth = normalize(ground_truth)
ground_truth = ground_truth.unsqueeze(dim=0)
bboxes = test_bbox(config,
batch_size=ground_truth.size(0),
t=50,
l=50)
x, mask = mask_image(ground_truth, bboxes, config)
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
g_checkpoint = torch.load(f'{checkpoint_path}/gen.pt')
netG.load_state_dict(g_checkpoint)
# model_iteration = int(last_model_name[-11:-3])
print("Model Resumed".format(checkpoint_path))
if cuda:
netG = nn.parallel.DataParallel(
netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2 = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result_cpu = torch.Tensor.cpu(
inpainted_result).detach().permute(0, 2, 3, 1)
inpainted_result_cpu = np.asarray(
inpainted_result_cpu[0])
inpainted_result_cpu = cv2.normalize(
inpainted_result_cpu, inpainted_result_cpu, 0, 255,
cv2.NORM_MINMAX)
# cat_result = torch.cat([x, inpainted_result, ground_truth], dim=3).cuda()
vutils.save_image(inpainted_result,
args.output_dir +
'output_{}/'.format(j + 1) +
'output_{}.png'.format(i),
padding=0,
normalize=True)
# cv2.imwrite(args.output_dir+ 'output_{}/'.format(j+1) + 'output_{}.png'.format(i), inpainted_result_cpu)
# cv2.cvtColor(inpainted_result_cpu, cv2.COLOR_BGR2RGB))
print("{}th image saved".format(i))
else:
raise TypeError("{} is not an image file.".format)
# exit no grad context
except Exception as e: # for unexpected error logging
print("Error: {}".format(e))
raise e
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
print("Configuration: {}".format(config))
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(args.image):
if args.mask and is_image_file(args.mask):
# Test a single masked image with a given mask
x = tif_loader(args.image)
## center crop
x = x[110:366, 110:366, :]
x = torch.from_numpy(x)
mask = default_loader(args.mask) ## 476 --> 256
#x = x[110:366, 110:366,:]
#x = transforms.Resize(config['image_shape'][:-1])(x)
#x = transforms.CenterCrop(config['image_shape'][:-1])(x)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(config['image_shape'][:-1])(mask)
#x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
x = normalize(x)
#x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
elif args.mask:
raise TypeError("{} is not an image file.".format(args.mask))
else:
# Test a single ground-truth image with a random mask
def __getitem__(self, index):
# first, load gt image
gt_path = os.path.join(self.gt_path, self.samples[index])
img = default_loader(gt_path, chan='L')
# get original image for GT images
# png_to_jpg_ext = self.samples[index]
original_path = os.path.join(self.data_path, self.samples[index])
original_path = original_path.replace('png', 'jpg')
# but, if the original image is not exist, skip it
orig_img = default_loader(original_path)
#img_path = os.path.join(self.data_path, self.gt_samples[index])
#print("====img(labeled)")
#print(img.size)
if self.random_crop:
# GT IMAGE
# GT IMAGE -> 128*128*3 -----> 128*128*9
imgw, imgh = img.size
if imgh < self.image_shape[0] or imgw < self.image_shape[1]:
img = transforms.Resize(min(self.image_shape))(img)
img = transforms.RandomCrop(self.image_shape)(img)
# ORIGINAL IMAGE
imgw, imgh = orig_img.size
if imgh < self.image_shape[0] or imgw < self.image_shape[1]:
orig_img = transforms.Resize(min(self.image_shape))(orig_img)
orig_img = transforms.RandomCrop(self.image_shape)(orig_img)
else:
img = transforms.Resize(self.image_shape)(img)
img = transforms.RandomCrop(self.image_shape)(img)
orig_img = transforms.Resize(self.image_shape)(orig_img)
orig_img = transforms.RandomCrop(self.image_shape)(orig_img)
# img : (0 ~ 15), we have to convert it to 16 channel
#img = np.array(img, dtype=np.int32)
img = np.array(img, dtype=np.uint8)
#print("max : " + str(img.max()) + ", low : " + str(img.min()))
#img = transforms.ToTensor()(img).int() # turn the image to a tensor
img = torch.from_numpy(img).long()
#print(">max : " + str(img.max()) + ", low : " + str(img.min()))
#\img = normalize(img)
orig_img = transforms.ToTensor()(
orig_img) # turn the image to a tensor
orig_img = normalize(orig_img)
raw_name = self.samples[index].split('.')[0]
# =============================== OUTPUT IMAGE CAUTION SIZE!!!
#target = np.zeros([self.n_classes, 128, 128])
target = np.zeros([self.n_classes, 256, 256])
for c in range(self.n_classes):
target[c][img == c] = 1
target = torch.from_numpy(target).float()
#lbl_img = m.toimage(img, high=img.max(), low=img.min())
#print(lbl_img.shape)
# RETURN (NAME), GT, TARGET(ONE-HOT), ORIG
if self.return_name:
return raw_name, img, target, orig_img
else:
return img, target, orig_img
