def __init__(self, left_image, right_image, left_disparity,
right_disparity, img_dir):
imgLeft = Image.open(left_image)
(self.col, self.row) = imgLeft.size
self.Left_Pixels = imgLeft.load()
imgLeft.close()
imgRight = Image.open(right_image)
self.Right_Pixels = imgRight.load()
imgRight.close()
self.MEGA_MATRIX_LEFT = np.zeros((self.row, self.col, self.num_sigmas,
self.num_thetas, self.num_channels))
self.MEGA_MATRIX_RIGHT = np.zeros((self.row, self.col, self.num_sigmas,
self.num_thetas, self.num_channels))
dispfile0 = open(left_disparity, 'rb')
self.arr0, scale0 = load_pfm.load_pfm(dispfile0)
dispfile1 = open(right_disparity, 'rb')
self.arr1, scale1 = load_pfm.load_pfm(dispfile1)
self.arr0 = np.flipud(self.arr0)
self.arr1 = np.flipud(self.arr1)
dispfile0.close()
dispfile1.close()
self.DIR = img_dir.rstrip('/') + "/"
def __getitem__(self, i):
"""
Get the ith item from the dataset
Return : left_img, right_img, target
"""
img_name = self.files[self.split][i]
left_img_path = self.root + self.split + self.Resolution + '/' + img_name+ '/' + 'im0.png'
right_img_path = self.root + self.split + self.Resolution + '/' + img_name+ '/' + 'im1_rectified.png'
lbl_path = self.root + "GT-" + self.split + self.Resolution + '/' + img_name+ '/' +'disp0GT.pfm'
left_img = m.imread(left_img_path)
left_img = np.array(left_img, dtype=np.float32)
left_img = left_img.transpose(2, 0, 1)
right_img = m.imread(right_img_path)
right_img = np.array(right_img, dtype=np.float32)
right_img = right_img.transpose(2, 0, 1)
# Normalizing images
left_img = (left_img - left_img.mean()) / left_img.std()
right_img = (right_img - right_img.mean()) / right_img.std()
lbl = load_pfm(lbl_path)
lbl = np.array(lbl[0], dtype=np.int64)
lbl[lbl==np.inf] = -1
if self._transform:
left_img, right_img, lbl = self.transform(img, lbl)
return left_img, right_img, lbl
def gen_datapoint(fname_disparity,
fname_disparity_next_frame,
fname_disparity_change,
fname_optical_flow,
image,
image_next_frame,
n=8192,
max_cut=35,
focal_length=1050.):
np.random.seed(0)
##### generate needed data
disparity_np, _ = load_pfm.load_pfm(fname_disparity)
disparity_next_frame_np, _ = load_pfm.load_pfm(fname_disparity_next_frame)
disparity_change_np, _ = load_pfm.load_pfm(fname_disparity_change)
optical_flow_np, _ = load_pfm.load_pfm(fname_optical_flow)
rgb_np = cv2.imread(image)[:, :, ::-1] / 255.
rgb_next_frame_np = cv2.imread(image_next_frame)[:, :, ::-1] / 255.
depth_np = focal_length / disparity_np
depth_next_frame_np = focal_length / disparity_next_frame_np
future_depth_np = focal_length / (disparity_np + disparity_change_np)
##### generate needed data
h, w = disparity_np.shape
##### point set 1 current pos
try:
depth_requirement = depth_np < max_cut
except:
return None
satisfy_pix1 = np.column_stack(np.where(depth_requirement))
if satisfy_pix1.shape[0] < n:
return None
sample_choice1 = np.random.choice(satisfy_pix1.shape[0],
size=n,
replace=False)
sampled_pix1_y = satisfy_pix1[sample_choice1, 0]
sampled_pix1_x = satisfy_pix1[sample_choice1, 1]
current_pos1 = np.array([
get_3d_pos_xy(sampled_pix1_y[i], sampled_pix1_x[i],
depth_np[int(sampled_pix1_y[i]),
int(sampled_pix1_x[i])]) for i in range(n)
])
current_rgb1 = np.array([[
rgb_np[h - 1 - int(sampled_pix1_y[i]),
int(sampled_pix1_x[i]),
0], rgb_np[h - 1 - int(sampled_pix1_y[i]),
int(sampled_pix1_x[i]), 1],
rgb_np[h - 1 - int(sampled_pix1_y[i]),
int(sampled_pix1_x[i]), 2]
] for i in range(n)])
##### point set 1 current pos
##### point set 1 future pos
sampled_optical_flow_x = np.array([
optical_flow_np[int(sampled_pix1_y[i]),
int(sampled_pix1_x[i])][0] for i in range(n)
])
sampled_optical_flow_y = np.array([
optical_flow_np[int(sampled_pix1_y[i]),
int(sampled_pix1_x[i])][1] for i in range(n)
])
future_pix1_x = sampled_pix1_x + sampled_optical_flow_x
future_pix1_y = sampled_pix1_y - sampled_optical_flow_y
future_pos1 = np.array([
get_3d_pos_xy(
future_pix1_y[i], future_pix1_x[i],
future_depth_np[int(sampled_pix1_y[i]),
int(sampled_pix1_x[i])]) for i in range(n)
])
##### point set 1 future pos
flow = future_pos1 - current_pos1
##### point set 2 current pos
try:
depth_requirement = depth_next_frame_np < max_cut
except:
return None
satisfy_pix2 = np.column_stack(np.where(depth_next_frame_np < max_cut))
if satisfy_pix2.shape[0] < n:
return None
sample_choice2 = np.random.choice(satisfy_pix2.shape[0],
size=n,
replace=False)
sampled_pix2_y = satisfy_pix2[sample_choice2, 0]
sampled_pix2_x = satisfy_pix2[sample_choice2, 1]
current_pos2 = np.array([
get_3d_pos_xy(
sampled_pix2_y[i], sampled_pix2_x[i],
depth_next_frame_np[int(sampled_pix2_y[i]),
int(sampled_pix2_x[i])]) for i in range(n)
])
current_rgb2 = np.array([[
rgb_next_frame_np[h - 1 - int(sampled_pix2_y[i]),
int(sampled_pix2_x[i]),
0], rgb_next_frame_np[h - 1 - int(sampled_pix2_y[i]),
int(sampled_pix2_x[i]), 1],
rgb_next_frame_np[h - 1 - int(sampled_pix2_y[i]),
int(sampled_pix2_x[i]), 2]
] for i in range(n)])
##### point set 2 current pos
##### mask, judge whether point move out of fov or occluded by other object after motion
future_pos1_depth = future_depth_np[sampled_pix1_y, sampled_pix1_x]
future_pos1_foreground_depth = np.zeros_like(future_pos1_depth)
valid_mask_fov1 = np.ones_like(future_pos1_depth, dtype=bool)
for i in range(future_pos1_depth.shape[0]):
if future_pix1_y[i] > 0 and future_pix1_y[i] < h and future_pix1_x[
i] > 0 and future_pix1_x[i] < w:
future_pos1_foreground_depth[i] = bilinear_interp_val(
depth_next_frame_np, future_pix1_y[i], future_pix1_x[i])
else:
valid_mask_fov1[i] = False
valid_mask_occ1 = (future_pos1_foreground_depth -
future_pos1_depth) > -5e-1
mask1 = valid_mask_occ1 & valid_mask_fov1
##### mask, judge whether point move out of fov or occluded by other object after motion
return current_pos1, current_pos2, current_rgb1, current_rgb2, flow, mask1