def visual_persp(img_glob, output_dir, device=torch.device("cpu"), path_prefix='ckpt/pre', flip=True, rotate=[], d1=21, d2=3): # Check input arguments validation for path in glob.glob(img_glob): assert os.path.isfile(path), '%s not found' % path assert os.path.isdir(output_dir), '%s is not a directory' % output_dir for rotate in rotate: assert 0 <= rotate and rotate <= 1, 'elements in --rotate should in [0, 1]' # Prepare model layoutnet = LayoutNet().to(device) torch_pretrained = torchfile.load( '/home/jupyter/Shapes.ai/pytorch-layoutnet/ckpt/perspfull_lsun_type_pretrained.t7' ) total_parameter = 0 for p in layoutnet.parameters(): total_parameter += np.prod(p.size()) print('pytorch parameters: ', total_parameter) print('t7 file: ', torch_pretrained.shape[0]) idx = 0 idx = copy_params(idx, layoutnet.parameters(), torch_pretrained) torch.save(layoutnet.state_dict(), '/ckpt/persp_pretrained.pth') #layoutnet.load_state_dict(torch.load('/ckpt/persp_pretrained.pth')) # Load path to visualization img_paths = sorted(glob.glob(img_glob)) # Process each input for i_path in img_paths: print('img path:', i_path) # Load input images i_img = np.array(Image.open(i_path).resize( (512, 512)), np.float32) / 255 #l_img = np.array(Image.open(l_path), np.float32) / 255 x_img = i_img.transpose([2, 0, 1]) print('x_img shape: ', x_img.shape) # Augment data x_imgs_augmented, aug_type = augment(x_img, flip, rotate) # Feedforward and extract output images with torch.no_grad(): x = torch.FloatTensor(x_imgs_augmented).to(device) print('x shape: ', x.shape) deconv6_sf, deconv6_sf_c, ref4 = layoutnet(x) #edg_de_list = edg_decoder(en_list[::-1]) #cor_de_list = cor_decoder(en_list[-1:] + edg_de_list[:-1]) edg_tensor = deconv6_sf cor_tensor = deconv6_sf_c roomtype_tensor = ref4 # Recover the effect from augmentation edg_img = augment_undo(edg_tensor.cpu().numpy(), aug_type) cor_img = augment_undo(cor_tensor.cpu().numpy(), aug_type) # Merge all results from augmentation edgmap = edg_img.transpose([0, 2, 3, 1]).mean(0).copy() cormap = cor_img.transpose([0, 2, 3, 1]).mean(0)[..., 0].copy() # Post processing to extract layout cor_id = get_ini_cor(cormap, d1, d2) #if post_optimization: # cor_id = optimize_cor_id(cor_id, edgmap, cormap, # num_iters=100, verbose=False) # Draw extracted layout on source image bon_img = draw_boundary_from_cor_id(cor_id.copy(), i_img * 255) # Composite all result in one image all_in_one = 0.3 * edgmap + 0.3 * cormap[..., None] + 0.4 * i_img all_in_one = draw_boundary_from_cor_id(cor_id.copy(), all_in_one * 255) # Dump results basename = os.path.splitext(os.path.basename(i_path))[0] path_edg = os.path.join(output_dir, '%s_edg.png' % basename) path_cor = os.path.join(output_dir, '%s_cor.png' % basename) path_bon = os.path.join(output_dir, '%s_bon.png' % basename) path_all_in_one = os.path.join(output_dir, '%s_all.png' % basename) path_cor_id = os.path.join(output_dir, '%s_cor_id.txt' % basename) Image.fromarray((edgmap * 255).astype(np.uint8)).save(path_edg) Image.fromarray((cormap * 255).astype(np.uint8)).save(path_cor) Image.fromarray(bon_img).save(path_bon) Image.fromarray(all_in_one).save(path_all_in_one) with open(path_cor_id, 'w') as f: for x, y in cor_id: f.write('%.6f %.6f\n' % (x, y))
# Merge all results from augmentation edgmap = edg_img.transpose([0, 2, 3, 1]).mean(0).copy() cormap = cor_img.transpose([0, 2, 3, 1]).mean(0)[..., 0].copy() # Post processing to extract layout cor_id = get_ini_cor(cormap, args.d1, args.d2) if args.post_optimization: cor_id = optimize_cor_id(cor_id, edgmap, cormap, num_iters=100, verbose=False) # Draw extracted layout on source image bon_img = draw_boundary_from_cor_id(cor_id.copy(), i_img * 255) # Composite all result in one image all_in_one = 0.3 * edgmap + 0.3 * cormap[..., None] + 0.4 * i_img all_in_one = draw_boundary_from_cor_id(cor_id.copy(), all_in_one * 255) # Dump results basename = os.path.splitext(os.path.basename(i_path))[0] path_edg = os.path.join(args.output_dir, '%s_edg.png' % basename) path_cor = os.path.join(args.output_dir, '%s_cor.png' % basename) path_bon = os.path.join(args.output_dir, '%s_bon.png' % basename) path_all_in_one = os.path.join(args.output_dir, '%s_all.png' % basename) path_cor_id = os.path.join(args.output_dir, '%s_cor_id.txt' % basename) Image.fromarray((edgmap * 255).astype(np.uint8)).save(path_edg) Image.fromarray((cormap * 255).astype(np.uint8)).save(path_cor)
def __getitem__(self, idx): pkl_path = self.root_dir + self.namelist[idx] pkl = pkl = pickle.load(open(pkl_path, 'rb')) img = pkl['image'].astype('float32') label = pkl['edge'].astype('float32') mask = pkl['line'].astype('float32') cor = pkl['cor'].astype('float32') # data augmentation if self.train_type == 'train': # random streching xmin, ymin, xmax, ymax = cor2xybound(cor) kx = np.random.uniform(1.0, self.max_stretch) ky = np.random.uniform(1.0, self.max_stretch) if np.random.randint(2) == 0: kx = max(1 / kx, min(0.5 / xmin, 1.0)) else: kx = min(kx, max(10.0 / xmax, 1.0)) if np.random.randint(2) == 0: ky = max(1 / ky, min(0.5 / ymin, 1.0)) else: ky = min(ky, max(10.0 / ymax, 1.0)) img, mask, cor = panostretch.pano_stretch(img, mask, cor, kx, ky) # random rotation random.seed() h = img.shape[0] w = img.shape[1] rot = int(np.floor(np.random.random() * w)) img = np.concatenate((img[:, rot:, :], img[:, :rot, :]), axis=1) mask = np.concatenate((mask[:, rot:, :], mask[:, :rot, :]), axis=1) cor[:, 0] = cor[:, 0] - rot id = cor[:, 0] < 0 cor[id, 0] = cor[id, 0] + 1023 # generate line label # sort gt cor_id = np.argsort(cor[:, 0]) cor = cor[cor_id, :] for row in range(0, cor.shape[0], 2): cor_id = np.argsort(cor[row:row + 2, 1]) cor[row:row + 2, :] = cor[row:row + 2, cor_id] # wall kpmap_w = np.zeros((self.im_h, self.im_w)) bg = np.zeros_like(img) for l_id in range(0, cor.shape[0], 2): panoEdgeC = draw_boundary_from_cor_id(cor[l_id:l_id + 2, :], bg) # add gaussian panoEdgeC = panoEdgeC.astype('float32') / 255.0 panoEdgeC = gaussian_filter(panoEdgeC[:, :, 1], sigma=20) panoEdgeC = panoEdgeC / np.max(panoEdgeC) kpmap_w = np.maximum(kpmap_w, panoEdgeC) # ceil kpmap_c = np.zeros((self.im_h, self.im_w)) cor_all = cor[[0, 2, 2, 4, 4, 6, 6, 0], :] for l_id in range(0, cor_all.shape[0], 2): panoEdgeC = draw_boundary_from_cor_id(cor_all[l_id:l_id + 2, :], bg) # add gaussian panoEdgeC = panoEdgeC[:, :, 1].astype('float32') / 255.0 panoEdgeC[int(np.amax(cor_all[l_id:l_id + 2, 1])) + 5:, :] = 0 panoEdgeC = gaussian_filter(panoEdgeC, sigma=20) panoEdgeC = panoEdgeC / np.max(panoEdgeC) kpmap_c = np.maximum(kpmap_c, panoEdgeC) # floor kpmap_f = np.zeros((self.im_h, self.im_w)) cor_all = cor[[1, 3, 3, 5, 5, 7, 7, 1], :] for l_id in range(0, cor_all.shape[0], 2): panoEdgeC = draw_boundary_from_cor_id(cor_all[l_id:l_id + 2, :], bg) # add gaussian panoEdgeC = panoEdgeC[:, :, 1].astype('float32') / 255.0 panoEdgeC[:int(np.amin(cor_all[l_id:l_id + 2, 1])) - 5, :] = 0 panoEdgeC = gaussian_filter(panoEdgeC, sigma=20) panoEdgeC = panoEdgeC / np.max(panoEdgeC) kpmap_f = np.maximum(kpmap_f, panoEdgeC) label = np.stack((kpmap_w, kpmap_c, kpmap_f), axis=-1) # generate corner label label2 = np.zeros((self.im_h, self.im_w)) for l_id in range(cor.shape[0]): panoEdgeC = np.zeros((self.im_h, self.im_w)) hh = int(np.round(cor[l_id, 1])) ww = int(np.round(cor[l_id, 0])) panoEdgeC[hh - 1:hh + 2, ww] = 1.0 panoEdgeC[hh, ww - 1:ww + 2] = 1.0 # add gaussian panoEdgeC = gaussian_filter(panoEdgeC, sigma=20) panoEdgeC = panoEdgeC / np.max(panoEdgeC) label2 = np.maximum(label2, panoEdgeC) label2 = np.expand_dims(label2, axis=2) if self.train_type == 'train': # gamma random.seed() g_prob = np.random.random() * 1 + 0.5 img = exposure.adjust_gamma(img, g_prob) # intensity random.seed() g_prob = np.random.random() * 127 img = exposure.rescale_intensity(img * 255.0, in_range=(g_prob, 255)) # color channel random.seed() g_prob = np.random.random() * 0.4 + 0.8 img[:, :, 0] = img[:, :, 0] * g_prob random.seed() g_prob = np.random.random() * 0.4 + 0.8 img[:, :, 1] = img[:, :, 1] * g_prob random.seed() g_prob = np.random.random() * 0.4 + 0.8 img[:, :, 2] = img[:, :, 2] * g_prob # random flip if random.uniform(0, 1) > 0.5: img = np.fliplr(img).copy() mask = np.fliplr(mask).copy() label = np.fliplr(label).copy() label2 = np.fliplr(label2).copy() # reshape np.clip(img, 0.0, 1.0, out=img) np.clip(label, 0.0, 1.0, out=label) np.clip(label2, 0.0, 1.0, out=label2) np.clip(mask, 0.0, 1.0, out=mask) img = np.concatenate((img, mask), axis=2) # permute dim if self.transform: if self.train_type == 'train': img = data_transforms['train'](img).float() label = data_transforms['train'](label).float() label2 = data_transforms['train'](label2).float() else: img = data_transforms['val'](img).float() label = data_transforms['val'](label).float() label2 = data_transforms['val'](label2).float() return img, label, label2