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