def inference(net, x, device, flip=False, rotate=[], visualize=False):
'''
net : the trained HorizonNet
x : tensor in shape [1, 3, 512, 1024]
flip : fliping testing augmentation
rotate: horizontal rotation testing augmentation
'''
H, W = tuple(x.shape[2:])
# Network feedforward (with testing augmentation)
x, aug_type = augment(x, args.flip, args.rotate)
y_bon_, y_cor_ = net(x.to(device))
y_bon_ = augment_undo(y_bon_.cpu(), aug_type).mean(0)
y_cor_ = augment_undo(torch.sigmoid(y_cor_).cpu(), aug_type).mean(0)
# Visualize raw model output
if visualize:
vis_out = visualize_a_data(x[0], torch.FloatTensor(y_bon_[0]),
torch.FloatTensor(y_cor_[0]))
else:
vis_out = None
y_bon_ = (y_bon_[0] / np.pi + 0.5) * H - 0.5
y_cor_ = y_cor_[0, 0]
# Detech wall-wall peaks (cuboid version)
xs_ = find_N_peaks(y_cor_, r=29, min_v=0, N=4)[0]
# Init floor/ceil plane
z0 = 50
_, z1 = post_proc.np_refine_by_fix_z(*y_bon_, z0)
# Generate cuboid wall-wall
cor, _ = post_proc.gen_ww(xs_,
y_bon_[0],
z0,
tol=abs(0.16 * z1 / 1.6),
force_cuboid=True)
# Expand with btn coory
cor = np.hstack(
[cor, post_proc.infer_coory(cor[:, 1], z1 - z0, z0)[:, None]])
assert len(cor) == 4, 'Extracted corner not cuboid !?'
# Collect corner position in equirectangular
cor_id = np.zeros((8, 2), np.float32)
for j in range(4):
cor_id[j * 2] = cor[j, 0], cor[j, 1]
cor_id[j * 2 + 1] = cor[j, 0], cor[j, 2]
# Normalized to [0, 1]
cor_id[:, 0] /= W
cor_id[:, 1] /= H
return cor_id, z0, z1, vis_out
def inference(net,
x,
device,
flip=False,
rotate=[],
visualize=False,
force_cuboid=False,
force_raw=False,
min_v=None,
r=0.05):
'''
net : the trained HorizonNet
x : tensor in shape [1, 3, 512, 1024]
flip : fliping testing augmentation
rotate: horizontal rotation testing augmentation
'''
H, W = tuple(x.shape[2:])
# Network feedforward (with testing augmentation)
x, aug_type = augment(x, flip, rotate)
y_bon_, y_cor_ = net(x.to(device))
y_bon_ = augment_undo(y_bon_.cpu(), aug_type).mean(0)
y_cor_ = augment_undo(torch.sigmoid(y_cor_).cpu(), aug_type).mean(0)
# Visualize raw model output
if visualize:
vis_out = visualize_a_data(x[0], torch.FloatTensor(y_bon_[0]),
torch.FloatTensor(y_cor_[0]))
else:
vis_out = None
y_bon_ = (y_bon_[0] / np.pi + 0.5) * H - 0.5
y_bon_[0] = np.clip(y_bon_[0], 1, H / 2 - 1)
y_bon_[1] = np.clip(y_bon_[1], H / 2 + 1, H - 2)
y_cor_ = y_cor_[0, 0]
# Init floor/ceil plane
z0 = 50
_, z1 = post_proc.np_refine_by_fix_z(*y_bon_, z0)
if force_raw:
# Do not run post-processing, export raw polygon (1024*2 vertices) instead.
# [TODO] Current post-processing lead to bad results on complex layout.
cor = np.stack([np.arange(1024), y_bon_[0]], 1)
else:
# Detech wall-wall peaks
if min_v is None:
min_v = 0 if force_cuboid else 0.05
r = int(round(W * r / 2))
N = 4 if force_cuboid else None
xs_ = find_N_peaks(y_cor_, r=r, min_v=min_v, N=N)[0]
# Generate wall-walls
cor, xy_cor = post_proc.gen_ww(xs_,
y_bon_[0],
z0,
tol=abs(0.16 * z1 / 1.6),
force_cuboid=force_cuboid)
if not force_cuboid:
# Check valid (for fear self-intersection)
xy2d = np.zeros((len(xy_cor), 2), np.float32)
for i in range(len(xy_cor)):
xy2d[i, xy_cor[i]['type']] = xy_cor[i]['val']
xy2d[i, xy_cor[i - 1]['type']] = xy_cor[i - 1]['val']
if not Polygon(xy2d).is_valid:
print(
'Fail to generate valid general layout!! '
'Generate cuboid as fallback.',
file=sys.stderr)
xs_ = find_N_peaks(y_cor_, r=r, min_v=0, N=4)[0]
cor, xy_cor = post_proc.gen_ww(xs_,
y_bon_[0],
z0,
tol=abs(0.16 * z1 / 1.6),
force_cuboid=True)
# Expand with btn coory
cor = np.hstack(
[cor, post_proc.infer_coory(cor[:, 1], z1 - z0, z0)[:, None]])
# Collect corner position in equirectangular
cor_id = np.zeros((len(cor) * 2, 2), np.float32)
for j in range(len(cor)):
cor_id[j * 2] = cor[j, 0], cor[j, 1]
cor_id[j * 2 + 1] = cor[j, 0], cor[j, 2]
# Normalized to [0, 1]
cor_id[:, 0] /= W
cor_id[:, 1] /= H
return cor_id, z0, z1, vis_out