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