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 test(dt_cor_id, z0, z1, gt_cor_id, w, h, losses):
# Eval corner error
mse = np.sqrt(((gt_cor_id - dt_cor_id)**2).sum(1)).mean()
ce_loss = 100 * mse / np.sqrt(w**2 + h**2)
# Pixel surface error (3 labels: ceiling, wall, floor)
y0_dt = []
y0_gt = []
y1_gt = []
for j in range(4):
coorxy = panostretch.pano_connect_points(dt_cor_id[j * 2],
dt_cor_id[(j * 2 + 2) % 8],
-z0)
y0_dt.append(coorxy)
coorxy = panostretch.pano_connect_points(gt_cor_id[j * 2],
gt_cor_id[(j * 2 + 2) % 8],
-z0)
y0_gt.append(coorxy)
coorxy = panostretch.pano_connect_points(gt_cor_id[j * 2 + 1],
gt_cor_id[(j * 2 + 3) % 8],
z0)
y1_gt.append(coorxy)
y0_dt = gen_reg_from_xy(np.concatenate(y0_dt, 0), w)
y1_dt = post_proc.infer_coory(y0_dt, z1 - z0, z0)
y0_gt = gen_reg_from_xy(np.concatenate(y0_gt, 0), w)
y1_gt = gen_reg_from_xy(np.concatenate(y1_gt, 0), w)
surface = np.zeros((h, w), dtype=np.int32)
surface[np.round(y0_dt).astype(int), np.arange(w)] = 1
surface[np.round(y1_dt).astype(int), np.arange(w)] = 1
surface = np.cumsum(surface, axis=0)
surface_gt = np.zeros((h, w), dtype=np.int32)
surface_gt[np.round(y0_gt).astype(int), np.arange(w)] = 1
surface_gt[np.round(y1_gt).astype(int), np.arange(w)] = 1
surface_gt = np.cumsum(surface_gt, axis=0)
pe_loss = 100 * (surface != surface_gt).sum() / (h * w)
# Eval 3d IoU
iou3d = eval_3diou(dt_cor_id[1::2], dt_cor_id[0::2], gt_cor_id[1::2],
gt_cor_id[0::2])
#print()
losses['CE'].append(ce_loss)
losses['PE'].append(pe_loss)
losses['3DIoU'].append(iou3d)
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
