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