def generate_triplets(self):
"""
Generate train test clusters
Sample multiple views per cluster in both train and test clusters
Iterate over train test clusters
Iterate over ref_view in sample_ref_views:
Generate Neighbors
Sample Neighbors
Iterate over sample_neighbors over generate_neighbors
Iterate over each bounding box
"""
if self.ref_views is None:
self.sample_ref_views()
proposals_per_img = self.sampling_params['proposals_per_img']
for ref_view in tqdm(self.ref_views, total=len(self.ref_views)):
self.logger.info(f'USING REFERENCE VIEW {ref_view}')
tc1w, Rc1w = get_tR(ref_view, self.image_struct)
if len(tc1w) == 0:
self.logger.debug(f'tR not available for img: {ref_view}')
print(f'tR not available for img: {ref_view}')
continue
tc1w *= self.scale
twc1, Rwc1 = camera_to_world_tR(tc1w, Rc1w)
bboxes = self.load_pt_file(ref_view)[:proposals_per_img]
if bboxes is None:
self.logger.exception(f'Proposal file: {ref_view} not found!!')
continue
img_path = os.path.join(self.scene_img_path, ref_view)
# img = Image.open(img_path).resize(self.image_params['proposal_size'])
img = Image.open(img_path)
depth_name = ref_view.split('.')[0][:-1] + '3.png'
depth_path = os.path.join(self.scene_depth_path, depth_name)
depth = Image.open(depth_path)
# ---------------------- SAVE OR PLOT PROPOSALS ------------------------
if self.plot_params['proposals'] or self.plot_params['save_proposals']:
fig, ax = bboxplot_in_img(img, bboxes, fontsize=7,
return_fig=True,
linewidth=2)
plt.title(f'Proposals for {ref_view}')
if self.plot_params['proposals']:
plt.show()
if self.plot_params['save_proposals']:
fig.savefig(os.path.join(self.plot_save_path,
f'proposals_{ref_view.split(".")[0]}.png'),
bbox_inches='tight', dpi=300)
plt.close()
# -----------------------------------------------------------------------
x, y, z = generate_flat_xyz(depth)
bboxes_px_idx = bbox_pixel_indices_list(np.array(bboxes),
x_flat=x,
y_flat=y,
z_flat=z,
filter_depth=False,
coordinates=False)
pcl_cam1, _ = project_xyz_to_camera(x_flat=x, y_flat=y, z_flat=z,
center_x=self.cx_prop, center_y=self.cy_prop,
focal_x=self.fx_prop, focal_y=self.fy_prop)
neighbors_dict = self.generate_neighbors(view=ref_view)
sampled_neighbors = self.sample_neighbors(neighbors_dict=neighbors_dict)
triplet_dict = OrderedDict()
for sampled_neighbor in sampled_neighbors:
triplet_dict[sampled_neighbor] = []
self.logger.info(f'\nCalculating for neighbor {sampled_neighbor}')
self.reproject_match_bboxes(ref_name=ref_view, ref_img=img, pcl_cam1=pcl_cam1,
ref_bboxes=bboxes, ref_bbox_indices=bboxes_px_idx,
neighbor_name=sampled_neighbor, twc1=twc1,
Rwc1=Rwc1,
output_list=triplet_dict[sampled_neighbor])
pickle_filename = ref_view.split('.')[0] + '.pickle'
pickle_path = os.path.join(self.triplet_save_path, pickle_filename)
with open(pickle_path, 'wb') as f:
pickle.dump(triplet_dict, f)
with open(os.path.join(self.output_scene_path, 'time.txt'), 'a') as f:
f.write("END TIME: " + datetime.now().strftime("%m-%d-%Y_%H:%M:%S") + '\n')
def reproject_match_bboxes(self, ref_name, ref_img, pcl_cam1, ref_bboxes, ref_bbox_indices,
neighbor_name, twc1,
Rwc1, output_list):
"""
Reproject all Camera World pixels of reference image to Neighbor View 2d
Load Neighbor Proposal bounding boxes
Skip this neighbor if proposal not available
For each bounding box in reference image
If num pixels inside reprojected image < threshold
Skip this bounding box
For each bounding box in neighbor image
Calculate iou
Find max iou
If max iou < threshold, skip
"""
tc2w, Rc2w = get_tR(neighbor_name, self.image_struct)
if len(tc2w) == 0:
self.logger.debug(f'tR not available for neighbor {neighbor_name}')
print(f'tR not available for neighbor {neighbor_name}')
return None
tc2w *= self.scale
tc2c1, Rc2c1 = inter_camera_tR(twc1, Rwc1, tc2w, Rc2w)
pcl_cam21 = np.matmul(Rc2c1, pcl_cam1) + tc2c1
proj21 = project_camera_to_2d(pcl_cam21, center_x=self.cx_prop, center_y=self.cy_prop,
focal_x=self.fx_prop, focal_y=self.fy_prop)
neighbor_bboxes = self.load_pt_file(neighbor_name)
if neighbor_bboxes is None:
# print(f'File not found maybe because no proposals in image!!')
self.logger.exception(
f'Proposal file: {neighbor_name} not found!!')
return None
elif len(neighbor_bboxes) == 0:
self.logger.debug(f'No bbox in {neighbor_name}')
return None
neighbor_bboxes = neighbor_bboxes[:self.sampling_params['proposals_per_neighbor_image']]
neighbor_img_path = os.path.join(self.scene_img_path, neighbor_name)
# neighbor_img = Image.open(neighbor_img_path).resize(self.image_params['proposal_size'])
neighbor_img = Image.open(neighbor_img_path)
for e, bbox_px_idx in enumerate(ref_bbox_indices):
current_ref_bbox = ref_bboxes[e]
proj_x = proj21[0][bbox_px_idx]
proj_y = proj21[1][bbox_px_idx]
proj_z = pcl_cam21[2][bbox_px_idx]
inside_image = np.logical_and.reduce(
(proj_x >= 0, proj_x < self.image_params['proposal_width'],
proj_y >= 0, proj_y < self.image_params['proposal_height'],
proj_z > 0))
valid_projected_pxl = np.sum(inside_image)
# FixMe: Include zero depths in count as well
total_projected_pxl = len(bbox_px_idx)
if valid_projected_pxl < self.match_params['min_valid_pixels']:
# print(f'Skipping valid pixels = {valid_projected_pxl}')
self.logger.debug(f'Skipping because valid pixels = {valid_projected_pxl}')
continue
proj_x = proj_x[inside_image]
proj_y = proj_y[inside_image]
matched_pixels_iou = []
for neighbor_bbox_idx, neighbor_bbox in enumerate(neighbor_bboxes):
area_bbox = (neighbor_bbox[2] - neighbor_bbox[0]) * (
neighbor_bbox[3] - neighbor_bbox[1])
proj_x_inside_bbox = np.logical_and(
proj_x >= neighbor_bbox[0],
proj_x <= neighbor_bbox[2])
proj_y_inside_bbox = np.logical_and(
proj_y >= neighbor_bbox[1],
proj_y <= neighbor_bbox[3])
num_inside_pixels = np.sum(
np.logical_and(proj_x_inside_bbox, proj_y_inside_bbox))
matched_pixels_iou.append(num_inside_pixels /
(area_bbox + (total_projected_pxl - num_inside_pixels)))
# TODO: Random sample for multiple same max_iou
max_iou_idx = np.argmax(matched_pixels_iou)
max_iou = matched_pixels_iou[max_iou_idx]
self.logger.info(f'Max IOU = {max_iou}')
pos_bbox = neighbor_bboxes[max_iou_idx]
# PLOT REF AND POS BBOX
# -------------------------------------------------------------------------
if self.plot_params['ref_pos'] or self.plot_params['save_ref_pos'] or \
self.plot_params['triplet'] or self.plot_params['save_triplet']:
fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(25, 14))
ax1.imshow(ref_img)
ax2.imshow(neighbor_img)
_ = scatterplot_in_img(neighbor_img, [proj_x, proj_y],
s=2, return_fig=True,
fig=fig,
ax=ax2)
_ = bboxplot_in_img(neighbor_img,
[pos_bbox],
fig=fig, ax=ax2,
numbering=False,
return_fig=True, linewidth=3)
_ = bboxplot_in_img(ref_img, [current_ref_bbox],
fig=fig, ax=ax1,
numbering=False,
return_fig=True, linewidth=3
)
ax1.title.set_text(ref_name)
ax2.title.set_text(neighbor_name)
ref_name_raw = ref_name.split(".")[0]
neighbor_name_raw = neighbor_name.split(".")[0]
fig.suptitle(f'{ref_name_raw}_{neighbor_name_raw}_{e}\nIOU = {max_iou:.3f}')
if self.plot_params['ref_pos']:
plt.show()
if self.plot_params['save_ref_pos']:
fig.savefig(os.path.join(self.plot_save_path,
f'{ref_name_raw}_{neighbor_name_raw}_{e}.png'),
dpi=300)
# -------------------------------------------------------------------------
if max_iou < self.match_params['min_pos_iou']:
self.logger.info(f'No bbox matched in neighbor {neighbor_name} because small IOU')
continue
neg_bbox = self.find_negative_bbox(bboxes=ref_bboxes, ref_bbox=current_ref_bbox)
if neg_bbox is None:
self.logger.debug(f'Could not find negative bbox!')
print(f'Could not find negative bbox!')
else:
output_list.append([current_ref_bbox, pos_bbox, neg_bbox])
# PLOT TRIPLET
# -------------------------------------------------------------------------
if self.plot_params['triplet'] or self.plot_params['save_triplet']:
# print('Triplet')
_ = bboxplot_in_img(ref_img,
[neg_bbox],
fig=fig, ax=ax1,
numbering=False, edgecolor='r',
return_fig=True, linewidth=3)
fig.suptitle(
f'Triplet_{ref_name_raw}_{neighbor_name_raw}_{e}\nIOU = {max_iou:.3f}')
if self.plot_params['triplet']:
# print('plot triplet')
# FixMe: Plot using plt because fig.show doesnt wait to close window
fig.show()
if self.plot_params['save_triplet']:
fig.savefig(os.path.join(self.plot_save_path,
f'Triplet_{ref_name_raw}_{neighbor_name_raw}_{e}.png'),
dpi=300)
plt.close()
# -------------------------------------------------------------------------
return None
def associate(self):
logger = logging.getLogger(__name__)
logging.basicConfig(filename=os.path.join('logs',
f'association_node_{time()}.log'),
filemode='w')
logger.setLevel(logging.DEBUG)
for zipped, triplet_folder in [[zip(self.train_centers, self.train_nodes,
self.train_anchor_views), 'train'], [zip(self.test_centers, self.test_nodes, self.test_anchor_views), 'test']]:
tqdm_length = len(self.train_nodes)
if triplet_folder == 'test':
tqdm_length = len(self.test_nodes)
for cluster_center, cluster_nodes, anchor_views in tqdm(zipped, total=tqdm_length):
proposal_filenames = [self.convert_jpg_pt(i) for i in anchor_views]
proposal_img_filenames = anchor_views
for proposal_filename, img_filename in zip(proposal_filenames,
proposal_img_filenames):
# print(f'Fetching samples for {img_filename}')
logger.info(f'Fetching samples for {img_filename}')
# print(f'Fetching samples for {img_filename}')
bboxes = self.load_pt_file(proposal_filename)
bboxes = bboxes[:self.proposals_per_img]
img_folder = self.scene
img_struct_folder = self.image_struct_dict[img_folder]
img_struct = img_struct_folder['image_struct']
img_scale = img_struct_folder['scale']
tc1w, Rc1w = get_tR(img_filename, img_struct)
if len(tc1w) == 0:
logger.debug(f'tR not available for img: {img_filename}')
print(f'tR not available for img: {img_filename}')
continue
tc1w *= img_scale
twc1, Rwc1 = camera_to_world_tR(tc1w, Rc1w)
img_path = os.path.join(self.dataset_root, img_folder,
'jpg_rgb',
img_filename)
img = Image.open(img_path).resize(image_defaults['size'])
camera_intrinsics = self.camera_intrinsics[img_folder]
if SAVE_FIG or PLOT_FIG:
fig, ax = bboxplot_in_img(img, bboxes, fontsize=10,
return_fig=True,
linewidth=3)
depth_name = self.convert_jpg_depth(img_filename)
depth_path = os.path.join(self.dataset_root, img_folder,
'high_res_depth', depth_name)
# Scale depth for new image size
depth = Image.open(depth_path).resize(image_defaults['size'],
resample=Image.NEAREST)
x, y, z = generate_flat_xyz(depth)
cx = camera_intrinsics['cx'] * x_resize_factor
cy = camera_intrinsics['cy'] * y_resize_factor
fx = camera_intrinsics['fx'] * x_resize_factor
fy = camera_intrinsics['fy'] * y_resize_factor
pcl_cam1, _ = project_xyz_to_camera(x_flat=x, y_flat=y,
z_flat=z,
center_x=cx,
center_y=cy, focal_x=fx,
focal_y=fy)
neighbor_names = self.sample_neighbor_nodes(node=cluster_center,
center_nodes=self.train_centers,
clusters=self.train_nodes)
neighbor_img_path = [
os.path.join(self.dataset_root, img_folder, 'jpg_rgb',
i) for i in neighbor_names]
neighbor_images = [Image.open(i).resize(image_defaults['size'])
for
i in neighbor_img_path]
logger.info(f'Sampled Neighbors : {neighbor_names}')
# print(f'Sampled Neighbors : {neighbor_names}')
bboxes_px_idx = bbox_pixel_indices_list(np.array(bboxes),
x_flat=x,
y_flat=y,
z_flat=z,
filter_depth=True,
coordinates=False)
if SAVE_FIG or PLOT_FIG:
for i in bboxes_px_idx:
fig, ax = scatterplot_in_img(img,
coordinates=(x[i], y[i]),
s=1,
fig=fig, ax=ax,
return_fig=True)
if SAVE_FIG:
plt.savefig(
os.path.join('/mnt/sda2/workspace/triplet_nodes_images',
img_filename.split('.')[0] + '.png'))
if PLOT_FIG:
plt.show()
plt.close()
# {Neighbor_name: [[1st triplet], ...]}
to_save = OrderedDict()
for neighbor_name, neighbor_img in zip(neighbor_names,
neighbor_images):
to_save[neighbor_name] = []
# print(f'\nCalculating for neighbor {neighbor_name}')
logger.info(f'\nCalculating for neighbor {neighbor_name}')
# print(f'\nCalculating for neighbor {neighbor_name}')
tc2w, Rc2w = get_tR(neighbor_name, img_struct)
if len(tc2w) == 0:
logger.debug(
f'tR not available for neighbor {neighbor_name}')
print(f'tR not available for neighbor {neighbor_name}')
continue
tc2w *= img_scale
tc2c1, Rc2c1 = inter_camera_tR(twc1, Rwc1, tc2w, Rc2w)
pcl_cam21 = np.matmul(Rc2c1, pcl_cam1) + tc2c1
proj21 = project_camera_to_2d(pcl_cam21, center_x=cx,
center_y=cy, focal_x=fx,
focal_y=fy)
neighbor_bboxes = self.load_pt_file(
self.convert_jpg_pt(neighbor_name))
if neighbor_bboxes is None:
# print(f'File not found maybe because no proposals in image!!')
logger.exception(
f'Proposal file: {neighbor_name} not found!!')
continue
neighbor_bboxes = neighbor_bboxes[
:self.proposals_neighbor_img]
# bboxplot_in_img(neighbor_img, neighbor_bboxes)
# Iterate over each of input image bbox
for e, bbox_px_idx in enumerate(bboxes_px_idx):
# Filter pixels outside image
# print(f'Bbox # = {e}')
logger.debug(f'Bbox # = {e}')
proj_x = proj21[0][bbox_px_idx]
proj_y = proj21[1][bbox_px_idx]
proj_z = pcl_cam21[2][bbox_px_idx]
# print(pcl_cam21.shape)
# scatterplot_in_img(neighbor_img,
# coordinates=(proj_x, proj_y),
# s=2)
inside_image = np.logical_and.reduce(
(proj_x >= 0, proj_x < image_defaults['width'],
proj_y >= 0, proj_y < image_defaults['height'],
proj_z > 0))
# Only for projected pixels with number of pixels>threshold
total_projected_pxl = len(bbox_px_idx)
valid_projected_pxl = np.sum(inside_image)
# TODO: Make this a parameter
# Skip if number of valid projected pixels is < threshold
if valid_projected_pxl < 15:
# print(f'Skipping valid pixels = {valid_projected_pxl}')
logger.debug(
f'Skipping valid pixels = {valid_projected_pxl}')
continue
proj_x = proj_x[inside_image]
proj_y = proj_y[inside_image]
# print(f'Num pixels = {valid_projected_pxl}')
logger.debug(f'Num pixels = {valid_projected_pxl}')
matched_pixels_iou = np.zeros(len(neighbor_bboxes))
# Iterate over each of SHUFFLED neighbor image bbox
neighbor_bbox_indices = np.random.choice(
list(range(len(neighbor_bboxes))),
size=len(neighbor_bboxes),
replace=False)
for neighbor_bbox_idx in neighbor_bbox_indices:
neighbor_bbox = neighbor_bboxes[neighbor_bbox_idx]
area = (neighbor_bbox[2] - neighbor_bbox[0]) * (
neighbor_bbox[3] - neighbor_bbox[1])
inside_proj_x = np.logical_and(
proj_x >= neighbor_bbox[0],
proj_x <= neighbor_bbox[2])
inside_proj_y = np.logical_and(
proj_y >= neighbor_bbox[1],
proj_y <= neighbor_bbox[3])
num_inside = np.sum(
np.logical_and(inside_proj_x, inside_proj_y))
# TODO: Note that only using area didn't produce good result because always prioritized smallest box with max pixels inside
matched_pixels_iou[neighbor_bbox_idx] = (
num_inside / (
area + (
total_projected_pxl - num_inside)))
matched_pixels = np.array(matched_pixels_iou)
matched_bbox_idx = matched_pixels.argmax()
# print(f'IOU = {matched_pixels[matched_bbox_idx]}')
logger.debug(
f'IOU = {matched_pixels[matched_bbox_idx]}')
# TODO: Make this a parameter
if matched_pixels[matched_bbox_idx] < 0.2:
logger.debug(f'Skipping because small IOU')
# print(f'Skipping because small IOU')
continue
if SAVE_FIG or PLOT_FIG:
fig, [ax2, ax1] = plt.subplots(1, 2,
figsize=(25, 14))
ax1.imshow(neighbor_img)
_ = scatterplot_in_img(neighbor_img,
[proj_x, proj_y],
s=2, return_fig=True,
fig=fig,
ax=ax1)
_ = bboxplot_in_img(neighbor_img,
[neighbor_bboxes[
matched_bbox_idx]],
fig=fig, ax=ax1,
numbering=False,
return_fig=True, linewidth=3)
current_bbox = bboxes[e]
found_neg = False
# TODO: Make this a parameter
max_search_count = min(10, len(bboxes))
search_idx_list = list(range(max_search_count))
# Shuffle the index list
np.random.shuffle(search_idx_list)
for neg_bbox_idx in search_idx_list:
neg_bbox = bboxes[neg_bbox_idx]
intersect_xmin = max(current_bbox[0], neg_bbox[0])
intersect_ymin = max(current_bbox[1], neg_bbox[1])
intersect_xmax = min(current_bbox[2], neg_bbox[2])
intersect_ymax = min(current_bbox[3], neg_bbox[3])
intersect_area = (
intersect_xmax - intersect_xmin) * (
intersect_ymax - intersect_ymin)
union_area = (current_bbox[2] - current_bbox[0]) * (
current_bbox[3] - current_bbox[1]) + (
neg_bbox[3] - neg_bbox[1]) * (
neg_bbox[2] - neg_bbox[
0]) - intersect_area
iou = intersect_area / union_area
if iou <= self.neg_bbox_iou_thresh:
# print(f'Neg bbox iou = {iou}')
logger.debug(f'Neg bbox iou = {iou}')
found_neg = True
break
if not found_neg:
logger.debug(
f'Did not find negative sample so skipping!!')
continue
if SAVE_FIG or PLOT_FIG:
ax2.imshow(img)
_ = bboxplot_in_img(img, [current_bbox, neg_bbox],
fig=fig, ax=ax2,
return_fig=True, linewidth=3)
if SAVE_FIG:
fig_path = os.path.join(
'/mnt/sda2/workspace/triplet_nodes_images',
img_filename.split('.')[0] + '_' +
neighbor_name.split('.')[
0] + '_' + str(e) + '.png'
)
plt.savefig(fig_path)
if PLOT_FIG:
plt.show()
plt.close()
to_save[neighbor_name].append(
[current_bbox, neighbor_bboxes[matched_bbox_idx],
neg_bbox])
pickle_filename = proposal_filename.split('.')[0] + '.pickle'
pickle_path = os.path.join(self.triplet_root, triplet_folder, pickle_filename)
with open(pickle_path, 'wb') as f:
pickle.dump(to_save, f)