def create_ceiling_floor_mask(cor_id, H, W):
'''
Binary masking on equirectangular
where 1 indicate floor or ceiling
'''
# Prepare 1d ceiling-wall/floor-wall boundary
c_pts = []
f_pts = []
n_cor = len(cor_id)
for i in range(n_cor // 2):
# Ceiling boundary points
xys = pano_connect_points(cor_id[i * 2],
cor_id[(i * 2 + 2) % n_cor],
z=-50,
w=W,
h=H)
c_pts.extend(xys)
# Floor boundary points
xys = pano_connect_points(cor_id[i * 2 + 1],
cor_id[(i * 2 + 3) % n_cor],
z=50,
w=W,
h=H)
f_pts.extend(xys)
# Sort for interpolate
c_pts = np.array(c_pts)
c_pts = c_pts[np.argsort(c_pts[:, 0] * H - c_pts[:, 1])]
f_pts = np.array(f_pts)
f_pts = f_pts[np.argsort(f_pts[:, 0] * H + f_pts[:, 1])]
# Removed duplicated point
c_pts = np.concatenate([c_pts[:1], c_pts[1:][np.diff(c_pts[:, 0]) > 0]], 0)
f_pts = np.concatenate([f_pts[:1], f_pts[1:][np.diff(f_pts[:, 0]) > 0]], 0)
# Generate boundary for each image column
c_bon = np.interp(np.arange(W), c_pts[:, 0], c_pts[:, 1])
f_bon = np.interp(np.arange(W), f_pts[:, 0], f_pts[:, 1])
# Generate mask
mask = np.zeros((H, W), np.bool)
for i in range(W):
u = max(0, int(round(c_bon[i])) + 1)
b = min(W, int(round(f_bon[i])))
mask[:u, i] = 1
mask[b:, i] = 1
return mask
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 cor_2_1d(cor, H, W):
bon_ceil_x, bon_ceil_y = [], []
bon_floor_x, bon_floor_y = [], []
n_cor = len(cor)
for i in range(n_cor // 2):
xys = panostretch.pano_connect_points(cor[i*2],
cor[(i*2+2) % n_cor],
z=-50, w=W, h=H)
bon_ceil_x.extend(xys[:, 0])
bon_ceil_y.extend(xys[:, 1])
for i in range(n_cor // 2):
xys = panostretch.pano_connect_points(cor[i*2+1],
cor[(i*2+3) % n_cor],
z=50, w=W, h=H)
bon_floor_x.extend(xys[:, 0])
bon_floor_y.extend(xys[:, 1])
bon_ceil_x, bon_ceil_y = sort_xy_filter_unique(bon_ceil_x, bon_ceil_y, y_small_first=True)
bon_floor_x, bon_floor_y = sort_xy_filter_unique(bon_floor_x, bon_floor_y, y_small_first=False)
bon = np.zeros((2, W))
bon[0] = np.interp(np.arange(W), bon_ceil_x, bon_ceil_y, period=W)
bon[1] = np.interp(np.arange(W), bon_floor_x, bon_floor_y, period=W)
bon = ((bon + 0.5) / H - 0.5) * np.pi
return bon
def __getitem__(self, idx):
# Read image
img_path = os.path.join(self.img_dir, self.img_fnames[idx])
img = np.array(Image.open(img_path), np.float32)[..., :3] / 255.
H, W = img.shape[:2]
# Read ground truth corners
with open(os.path.join(self.cor_dir, self.txt_fnames[idx])) as f:
cor = np.array(
[line.strip().split() for line in f if line.strip()],
np.float32)
# Corner with minimum x should at the beginning
cor = np.roll(cor[:, :2], -2 * np.argmin(cor[::2, 0]), 0)
# Detect occlusion
occlusion = find_occlusion(cor[::2].copy()).repeat(2)
assert (np.abs(cor[0::2, 0] - cor[1::2, 0]) >
W / 100).sum() == 0, img_path
assert (cor[0::2, 1] > cor[1::2, 1]).sum() == 0, img_path
# Stretch augmentation
if self.stretch:
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, cor = panostretch.pano_stretch(img, cor, kx, ky)
# Prepare 1d ceiling-wall/floor-wall boundary
bon_ceil_x, bon_ceil_y = [], []
bon_floor_x, bon_floor_y = [], []
n_cor = len(cor)
for i in range(n_cor // 2):
xys = panostretch.pano_connect_points(cor[i * 2],
cor[(i * 2 + 2) % n_cor],
z=-50)
bon_ceil_x.extend(xys[:, 0])
bon_ceil_y.extend(xys[:, 1])
for i in range(n_cor // 2):
xys = panostretch.pano_connect_points(cor[i * 2 + 1],
cor[(i * 2 + 3) % n_cor],
z=50)
bon_floor_x.extend(xys[:, 0])
bon_floor_y.extend(xys[:, 1])
bon_ceil_x, bon_ceil_y = sort_xy_filter_unique(bon_ceil_x,
bon_ceil_y,
y_small_first=True)
bon_floor_x, bon_floor_y = sort_xy_filter_unique(bon_floor_x,
bon_floor_y,
y_small_first=False)
bon = np.zeros((2, W))
bon[0] = np.interp(np.arange(W), bon_ceil_x, bon_ceil_y, period=W)
bon[1] = np.interp(np.arange(W), bon_floor_x, bon_floor_y, period=W)
bon = ((bon + 0.5) / img.shape[0] - 0.5) * np.pi
# Random flip
if self.flip and np.random.randint(2) == 0:
img = np.flip(img, axis=1)
bon = np.flip(bon, axis=1)
cor[:, 0] = img.shape[1] - 1 - cor[:, 0]
# Random horizontal rotate
if self.rotate:
dx = np.random.randint(img.shape[1])
img = np.roll(img, dx, axis=1)
bon = np.roll(bon, dx, axis=1)
cor[:, 0] = (cor[:, 0] + dx) % img.shape[1]
# Random gamma augmentation
if self.gamma:
p = np.random.uniform(1, 2)
if np.random.randint(2) == 0:
p = 1 / p
img = img**p
# Prepare 1d wall-wall probability
corx = cor[~occlusion, 0]
dist_o = cdist(corx.reshape(-1, 1),
np.arange(img.shape[1]).reshape(-1, 1),
p=1)
dist_r = cdist(corx.reshape(-1, 1),
np.arange(img.shape[1]).reshape(-1, 1) + img.shape[1],
p=1)
dist_l = cdist(corx.reshape(-1, 1),
np.arange(img.shape[1]).reshape(-1, 1) - img.shape[1],
p=1)
dist = np.min([dist_o, dist_r, dist_l], 0)
nearest_dist = dist.min(0)
y_cor = (self.p_base**nearest_dist).reshape(1, -1)
# Convert all data to tensor
x = torch.FloatTensor(img.transpose([2, 0, 1]).copy())
bon = torch.FloatTensor(bon.copy())
y_cor = torch.FloatTensor(y_cor.copy())
# Check whether additional output are requested
out_lst = [x, bon, y_cor]
if self.return_cor:
out_lst.append(cor)
if self.return_path:
out_lst.append(img_path)
return out_lst
def __getitem__(self, idx):
# Read image
img_path = os.path.join(self.img_dir, self.img_fnames[idx])
img = np.array(Image.open(img_path), np.float32)[..., :3] / 255.
H, W = img.shape[:2]
# init random variable
rnd_y, rnd_x = 0, 0
# Read ground truth corners
with open(os.path.join(self.cor_dir, self.txt_fnames[idx])) as f:
cor = np.array([line.strip().split() for line in f], np.float32)
# Corner with minimum x should at the beginning
cor = np.roll(cor[:, :2], -2 * np.argmin(cor[::2, 0]), 0)
# Detect occlusion
occlusion = find_occlusion(cor[::2].copy()).repeat(2)
assert (cor[0::2, 0] != cor[1::2, 0]).sum() == 0
assert (cor[0::2, 1] > cor[1::2, 1]).sum() == 0
# Stretch augmentation
if self.stretch:
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, cor = panostretch.pano_stretch(img, cor, kx, ky)
# Prepare 1d ceiling-wall/floor-wall boundary
bon = np.zeros((2, W))
bon[0, :] = 1e9
bon[1, :] = -1e9
n_cor = len(cor)
for i in range(n_cor // 2):
xys = panostretch.pano_connect_points(cor[i * 2],
cor[(i * 2 + 2) % n_cor],
z=-50)
xys = xys.astype(int)
bon[0, xys[:, 0]] = np.minimum(bon[0, xys[:, 0]], xys[:, 1])
for i in range(n_cor // 2):
xys = panostretch.pano_connect_points(cor[i * 2 + 1],
cor[(i * 2 + 3) % n_cor],
z=50)
xys = xys.astype(int)
bon[1, xys[:, 0]] = np.maximum(bon[1, xys[:, 0]], xys[:, 1])
bon = ((bon + 0.5) / img.shape[0] - 0.5) * np.pi
#print("boundary before {}".format(bon))
#print("corner before {}".format(cor))
# bon to relative height v
# Random flip
if self.flip and np.random.randint(2) == 0:
img = np.flip(img, axis=1)
bon = np.flip(bon, axis=1)
cor[:, 0] = img.shape[1] - 1 - cor[:, 0]
# Random horizontal rotate
if self.rotate:
dx = np.random.randint(img.shape[1])
img = np.roll(img, dx, axis=1)
bon = np.roll(bon, dx, axis=1)
cor[:, 0] = (cor[:, 0] + dx) % img.shape[1]
# Random gamma augmentation
if self.gamma:
p = np.random.uniform(1, 2)
if np.random.randint(2) == 0:
p = 1 / p
img = img**p
#random x rotate around x axis?
if (self.max_x_rotate > 0 and self.max_x_rotate > self.min_x_rotate):
rnd = np.random.uniform()
if (rnd > self.x_rotate_prob):
rnd_x = np.random.randint(self.min_x_rotate, self.max_x_rotate)
img = pano_lsd_align.rotatePanorama_degree(img, rnd_x, axis=1)
#bon,cor = pano_lsd_align.rotateCorners(bon,cor,rnd_x,axis=1) #simple :> its not working :< use -rnd_x instead of rnd_x here
#random z rotate around y axis?
if (self.max_y_rotate > 0 and self.max_y_rotate > self.min_y_rotate):
rnd = np.random.uniform()
if (rnd > self.y_rotate_prob):
rnd_y = np.random.randint(self.min_y_rotate, self.max_y_rotate)
img = pano_lsd_align.rotatePanorama_degree(img, rnd_y, axis=2)
#bon,cor = pano_lsd_align.rotateCorners(bon,cor,rnd_y,axis=2) #simple :> use -rnd_y instead of rnd_y here
#print("boundary after {}".format(bon))
#print("corner after {}".format(cor))
# Prepare 1d wall-wall probability
corx = cor[~occlusion, 0]
dist_o = cdist(corx.reshape(-1, 1),
np.arange(img.shape[1]).reshape(-1, 1),
p=1)
dist_r = cdist(corx.reshape(-1, 1),
np.arange(img.shape[1]).reshape(-1, 1) + img.shape[1],
p=1)
dist_l = cdist(corx.reshape(-1, 1),
np.arange(img.shape[1]).reshape(-1, 1) - img.shape[1],
p=1)
dist = np.min([dist_o, dist_r, dist_l], 0)
nearest_dist = dist.min(0)
y_cor = (self.p_base**nearest_dist).reshape(1, -1)
# Convert all data to tensor
x = torch.FloatTensor(img.transpose([2, 0, 1]).copy())
#print(x.shape)
bon = torch.FloatTensor(bon.copy())
y_cor = torch.FloatTensor(y_cor.copy())
#print("boundary final {}".format(bon))
#print("corner final {}".format(y_cor))
# Check whether additional output are requested
#rnd_x rnd_y are in degrees
out_lst = [x, bon, y_cor, rnd_x, rnd_y]
if self.return_cor:
out_lst.append(cor)
if self.return_path:
out_lst.append(img_path)
out_lst
return out_lst
