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