def load_from_dir(cls,
img_dir: str,
json_dir: str,
show_pbar: bool = True) -> NDDS_Frame_Handler:
check_dir_exists(json_dir)
check_dir_exists(img_dir)
img_pathlist = get_valid_image_paths(img_dir)
json_path_list = [
path for path in get_all_files_of_extension(dir_path=json_dir,
extension='json')
if not get_filename(path).startswith('_')
]
json_path_list.sort()
handler = NDDS_Frame_Handler()
if show_pbar:
pbar = tqdm(total=len(json_path_list), unit='ann(s)', leave=True)
pbar.set_description(f'Loading {cls.__name__}')
for json_path in json_path_list:
check_file_exists(json_path)
json_rootname = get_rootname_from_path(json_path)
matching_img_path = None
matching_cs_img_path = None
matching_depth_img_path = None
matching_is_img_path = None
for img_path in img_pathlist:
img_rootname = '.'.join(get_filename(img_path).split('.')[:-1])
if img_rootname == json_rootname:
matching_img_path = img_path
elif img_rootname == f'{json_rootname}.cs':
matching_cs_img_path = img_path
elif img_rootname == f'{json_rootname}.depth':
matching_depth_img_path = img_path
elif img_rootname == f'{json_rootname}.is':
matching_is_img_path = img_path
if matching_img_path and matching_cs_img_path and matching_depth_img_path and matching_is_img_path:
break
if matching_img_path is None:
logger.error(
f"Couldn't find image file that matches rootname of {get_filename(json_path)} in {img_dir}"
)
raise FileNotFoundError
frame = NDDS_Frame(
img_path=matching_img_path,
ndds_ann=NDDS_Annotation.load_from_path(json_path),
cs_img_path=matching_cs_img_path,
depth_img_path=matching_depth_img_path,
is_img_path=matching_is_img_path)
handler.append(frame)
if show_pbar:
pbar.update()
return handler
def _check_paths_valid(self, src_img_dir: str):
check_dir_exists(src_img_dir)
img_filename_list = []
duplicate_img_filename_list = []
for frame in self:
img_filename = get_filename(frame.img_path)
if img_filename not in img_filename_list:
img_filename_list.append(frame.img_path)
else:
duplicate_img_filename_list.append(frame.img_path)
img_path = f'{src_img_dir}/{img_filename}'
check_file_exists(img_path)
if frame.cs_img_path:
check_file_exists(
f'{src_img_dir}/{get_filename(frame.cs_img_path)}')
if frame.depth_img_path:
check_file_exists(
f'{src_img_dir}/{get_filename(frame.depth_img_path)}')
if frame.is_img_path:
check_file_exists(
f'{src_img_dir}/{get_filename(frame.is_img_path)}')
if len(duplicate_img_filename_list) > 0:
logger.error(
f'Found the following duplicate image filenames in {self.__class__.__name__}:\n{duplicate_img_filename_list}'
)
raise Exception
def update_images(
self, target_buffer: COCO_Field_Buffer, target_mapper: COCO_Mapper_Handler,
coco_image: COCO_Image, unique_key: str, new_img_filename: str=None, verbose: bool=False
):
pending_coco_image = coco_image.copy()
if new_img_filename is not None:
coco_url = pending_coco_image.coco_url
file_name = pending_coco_image.file_name
coco_url_dirpath = get_dirpath_from_filepath(coco_url)
coco_url_filename = get_filename(coco_url)
if file_name != coco_url_filename:
logger.error(f"file_name == {file_name} != coco_url_filename == {coco_url_filename}")
raise Exception
new_coco_url = f"{coco_url_dirpath}/{new_img_filename}"
new_file_name = new_img_filename
pending_coco_image.coco_url = new_coco_url
pending_coco_image.file_name = new_file_name
added_new, old_id, new_id = target_buffer.process_image(
coco_image=pending_coco_image, id_mapper=target_mapper, unique_key=unique_key
)
if verbose:
print_id_update(
label=f'{unique_key} Image', added_new=added_new,
old_id=old_id, new_id=new_id
)
def from_img_path(self, img_path: str, license_id: int,
image_id: int) -> COCO_Image:
check_file_exists(img_path)
img = cv2.imread(img_path)
img_h, img_w = img.shape[:2]
return COCO_Image(license_id=license_id,
file_name=get_filename(img_path),
coco_url=img_path,
height=img_h,
width=img_w,
date_captured=get_ctime(img_path),
flickr_url=None,
id=image_id)
def _check_paths_valid(self, src_img_dir: str):
check_dir_exists(src_img_dir)
img_filename_list = []
duplicate_img_filename_list = []
for ann in self:
img_filename = get_filename(ann.img_path)
if img_filename not in img_filename_list:
img_filename_list.append(ann.img_path)
else:
duplicate_img_filename_list.append(ann.img_path)
img_path = f'{src_img_dir}/{img_filename}'
check_file_exists(img_path)
if len(duplicate_img_filename_list) > 0:
logger.error(f'Found the following duplicate image filenames in LabelmeAnnotationHandler:\n{duplicate_img_filename_list}')
raise Exception
def load_split(self, preserve_filenames: bool=False, verbose: bool=False):
self.setup_directories(verbose=verbose)
random.shuffle(self.buffer.images.image_list)
train_image_list, test_image_list, val_image_list = self.split_image_list(image_list=self.buffer.images.image_list)
for image_list, target_dir, target_buffer, target_mapper, unique_key in \
zip(
[train_image_list, test_image_list, val_image_list],
[self.train_dir, self.test_dir, self.val_dir],
[self.train_buffer, self.test_buffer, self.val_buffer],
[self.buffer2train_mapper, self.buffer2test_mapper, self.buffer2val_mapper],
['train', 'test', 'val']
):
dest_img_dir = f"{target_dir}/img"
for coco_image in image_list:
image_id = coco_image.id
if not preserve_filenames:
img_extension = get_extension_from_filename(coco_image.file_name)
new_img_filename = get_filename(get_next_dump_path(dump_dir=dest_img_dir, file_extension=img_extension))
self.copy_image(coco_image=coco_image, dest_img_dir=dest_img_dir, new_img_filename=new_img_filename, verbose=verbose)
else:
self.copy_image(coco_image=coco_image, dest_img_dir=dest_img_dir, verbose=verbose)
coco_license = self.buffer.licenses.get_license_from_id(id=coco_image.license_id)
self.update_licenses(
target_buffer=target_buffer, target_mapper=target_mapper,
coco_license=coco_license, unique_key=unique_key, verbose=verbose
)
if not preserve_filenames:
self.update_images(
target_buffer=target_buffer, target_mapper=target_mapper,
coco_image=coco_image, unique_key=unique_key, new_img_filename=new_img_filename, verbose=verbose
)
else:
self.update_images(
target_buffer=target_buffer, target_mapper=target_mapper,
coco_image=coco_image, unique_key=unique_key, verbose=verbose
)
coco_annotation_list = self.buffer.annotations.get_annotations_from_imgIds(imgIds=[image_id])
for coco_annotation in coco_annotation_list:
coco_category = self.buffer.categories.get_category_from_id(id=coco_annotation.category_id)
self.update_categories(
target_buffer=target_buffer, target_mapper=target_mapper,
coco_category=coco_category, unique_key=unique_key, verbose=verbose
)
self.update_annotations(
target_buffer=target_buffer, target_mapper=target_mapper,
coco_annotation=coco_annotation, unique_key=unique_key, verbose=verbose
)
def get_frame_from_img_filename(self, img_filename: str) -> NDDS_Frame:
possible_filenames = []
for frame in self:
frame_filenames = [
get_filename(path) for path in [
frame.img_path, frame.is_img_path, frame.cs_img_path,
frame.depth_img_path
]
]
if img_filename in frame_filenames:
return frame
else:
possible_filenames.extend(frame_filenames)
logger.error(
f'Unable to find any image filename by the name of {img_filename} in NDDS_Frame_Handler.'
)
logger.error(f'Possible filenames:\n{possible_filenames}')
raise Exception
def copy_image(self, coco_image: COCO_Image, dest_img_dir: str, new_img_filename: str=None, verbose: bool=False):
coco_url = coco_image.coco_url
found, src_img_path = find_img_path(coco_url=coco_url, src_root_dir=self.src_root_dir)
if not found:
logger.error(f"Couldn't find any permutation of coco_url={coco_url} under {self.src_root_dir}")
raise Exception
if new_img_filename is None:
img_filename = get_filename(path=coco_url)
if img_filename != coco_image.file_name:
logger.error(f"coco_url filename doesn't match coco json file_name")
logger.error(f"coco_url: {coco_url}")
logger.error(f"file_name: {coco_image.file_name}")
raise Exception
else:
img_filename = new_img_filename
dest_img_path = f"{dest_img_dir}/{img_filename}"
silent = not verbose
copy_file(src_path=src_img_path, dest_path=dest_img_path, silent=silent)
def process_image(self,
coco_image: COCO_Image,
id_mapper: COCO_Mapper_Handler,
unique_key: str,
img_dir: str = None,
update_img_path: bool = False) -> (bool, int, int):
pending_coco_image = coco_image.copy()
if update_img_path:
if img_dir is None:
logger.error(
f"img_dir is required in order to update the img_path")
raise Exception
coco_url_filename = get_filename(pending_coco_image.coco_url)
if coco_url_filename != pending_coco_image.file_name:
logger.error(
f"coco_url_filename == {coco_url_filename} != pending_coco_image.file_name == {pending_coco_image.file_name}"
)
raise Exception
pending_coco_image.coco_url = f"{img_dir}/{coco_url_filename}"
# No checks should be necessary
old_id = pending_coco_image.id
new_id = len(self.images.image_list)
# Get License Id
found, new_license_id = id_mapper.license_mapper.get_new_id(
unique_key=unique_key, old_id=pending_coco_image.license_id)
if not found:
logger.error(
f"Couldn't find unique_key: {unique_key}, old_id: {pending_coco_image.license_id} pair in license_mapper."
)
raise Exception
# Update Id References
pending_coco_image.license_id = new_license_id
pending_coco_image.id = new_id
self.images.add(pending_coco_image)
added_new = True
id_mapper.image_mapper.add(unique_key=unique_key,
old_id=old_id,
new_id=new_id)
return added_new, old_id, new_id
def infer_coco_dataset(
self,
dataset: COCO_Dataset,
kpt_3d: np.ndarray,
corner_3d: np.ndarray,
K: np.ndarray,
camera_translation: np.ndarray = None,
camera_quaternion: np.ndarray = None,
distortion: np.ndarray = None,
line_start_point3d: np.ndarray = None,
line_end_point_3d: np.ndarray = None,
units_per_meter: float = 1.0,
blackout: bool = False,
dsize: (int, int) = None,
show_pbar: bool = True,
leave_pbar: bool = False,
preserve_dump_img_filename: bool = True,
accumulate_pred_dump: bool = True,
pred_dump_path: str = None,
test_name: str = None,
model_name: str = None,
stream_writer: StreamWriter = None,
leave_stream_writer_open: bool = False,
) -> PVNetFrameResultList:
pbar = (tqdm(total=len(dataset.images),
unit="image(s)",
leave=leave_pbar) if show_pbar else None)
frame_result_list = (PVNetFrameResultList() if pred_dump_path
is not None or accumulate_pred_dump else None)
for coco_image in dataset.images:
file_name = get_filename(coco_image.file_name)
if pbar is not None:
pbar.set_description(file_name)
img = Image.open(coco_image.coco_url)
orig_img_w, orig_img_h = img.size
if dsize is not None:
img = img.resize(dsize)
img_w, img_h = img.size
xscale = img_w / orig_img_w
yscale = img_h / orig_img_h
orig_img = np.asarray(img)
orig_img = cv2.cvtColor(orig_img, cv2.COLOR_RGB2BGR)
anns = dataset.annotations.get_annotations_from_imgIds(
[coco_image.id])
pred_list = PnpPredictionList()
for ann in anns:
working_img = orig_img.copy()
if dsize is not None:
bbox = ann.bbox.resize(
orig_frame_shape=[orig_img_h, orig_img_w, 3],
new_frame_shape=[img_h, img_w, 3],
)
else:
bbox = ann.bbox.copy()
if blackout:
bbox = bbox.clip_at_bounds(frame_shape=working_img.shape)
working_img = bbox.crop_and_paste(
src_img=working_img,
dst_img=np.zeros_like(working_img))
pred = self.predict(img=working_img,
bbox=bbox if blackout else None)
pred_list.append(
pred.to_pnp_pred(
gt_kpt_3d=kpt_3d,
corner_3d=corner_3d,
K=K,
camera_translation=camera_translation,
camera_quaternion=camera_quaternion,
distortion=distortion,
line_start_point3d=line_start_point3d,
line_end_point_3d=line_end_point_3d,
units_per_meter=units_per_meter,
))
if frame_result_list is not None:
frame_result = PVNetFrameResult(
frame=file_name,
pred_list=pred_list,
test_name=test_name,
model_name=model_name,
)
frame_result_list.append(frame_result)
result = pred_list.draw(img=orig_img)
if stream_writer is not None:
stream_writer.step(
img=result,
file_name=file_name
if preserve_dump_img_filename else None,
)
if pbar is not None:
pbar.update()
if stream_writer is not None and not leave_stream_writer_open:
stream_writer.close()
if pred_dump_path is not None:
frame_result_list.save_to_path(pred_dump_path, overwrite=True)
if pbar is not None:
pbar.close()
return frame_result_list
def infer_linemod_dataset(
self,
dataset: Linemod_Dataset,
img_dir: str,
blackout: bool = False,
show_pbar: bool = True,
show_preview: bool = False,
video_save_path: str = None,
dump_dir: str = None,
accumulate_pred_dump: bool = True,
pred_dump_path: str = None,
test_name: str = None,
) -> PVNetFrameResultList:
stream_writer = StreamWriter(
show_preview=show_preview,
video_save_path=video_save_path,
dump_dir=dump_dir,
)
pbar = tqdm(total=len(dataset.images), unit="image(s)") \
if show_pbar else None
frame_result_list = (PVNetFrameResultList() if pred_dump_path
is not None or accumulate_pred_dump else None)
for linemod_image in dataset.images:
file_name = get_filename(linemod_image.file_name)
if pbar is not None:
pbar.set_description(file_name)
img_path = f"{img_dir}/{file_name}"
img = Image.open(img_path)
orig_img = np.asarray(img)
orig_img = cv2.cvtColor(orig_img, cv2.COLOR_RGB2BGR)
result = orig_img.copy()
gt_ann = dataset.annotations.get(image_id=linemod_image.id)[0]
if blackout:
xmin = int(
gt_ann.corner_2d.to_numpy(demarcation=True)[:, 0].min())
xmax = int(
gt_ann.corner_2d.to_numpy(demarcation=True)[:, 0].max())
ymin = int(
gt_ann.corner_2d.to_numpy(demarcation=True)[:, 1].min())
ymax = int(
gt_ann.corner_2d.to_numpy(demarcation=True)[:, 1].max())
bbox = BBox(xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax)
bbox = bbox.clip_at_bounds(frame_shape=result.shape)
result = bbox.crop_and_paste(src_img=result,
dst_img=np.zeros_like(result))
else:
bbox = None
pred = self.predict(img=img, bbox=bbox)
if frame_result_list is not None:
frame_result = PVNetFrameResult(
frame=file_name,
pred_list=PVNetPredictionList([pred]),
test_name=test_name,
)
frame_result_list.append(frame_result)
gt_kpt_3d = gt_ann.fps_3d.copy()
gt_kpt_3d.append(gt_ann.center_3d)
pose_pred = pred.to_pose_pred(gt_kpt_3d=gt_kpt_3d, K=gt_ann.K)
corner_2d_gt = pvnet_pose_utils.project(
gt_ann.corner_3d.to_numpy(),
gt_ann.K.to_matrix(),
np.array(gt_ann.pose.to_list()),
)
corner_2d_pred = pred.to_corner_2d_pred(
gt_corner_3d=gt_ann.corner_3d, K=gt_ann.K, pose_pred=pose_pred)
result = draw_pts2d(img=result,
pts2d=gt_ann.fps_2d.to_numpy(),
color=(0, 255, 0),
radius=3)
result = draw_corners(img=result,
corner_2d=corner_2d_gt,
color=(0, 255, 0),
thickness=2)
result = draw_pts2d(img=result,
pts2d=pred.kpt_2d,
color=(0, 0, 255),
radius=2)
result = draw_corners(img=result,
corner_2d=corner_2d_pred,
color=(0, 0, 255),
thickness=2)
stream_writer.step(img=result, file_name=file_name)
if pbar is not None:
pbar.update()
stream_writer.close()
if pred_dump_path is not None:
frame_result_list.save_to_path(pred_dump_path, overwrite=True)
pbar.close()
return frame_result_list
def infer(
path: str,
weights_path: str,
thresh: int = 0.5,
key: str = 'R',
infer_dump_dir: str = '',
model: str = 'mask_rcnn_R_50_FPN_1x',
size: int = 1024,
class_names: List[str] = ['hook'],
gt_path:
str = '/home/jitesh/3d/data/coco_data/hook_test/json/cropped_hook.json'):
# class_names=['hook', 'pole']
# class_names=['hook']
conf_thresh = 0.001
show_bbox_border = True
gt_dataset = COCO_Dataset.load_from_path(json_path=gt_path)
inferer_seg = inferer(
weights_path=weights_path,
confidence_threshold=0.1,
# num_classes=1,
# num_classes=2,
class_names=class_names,
# class_names=['hook'],
model='keypoint_rcnn_R_50_FPN_1x',
# model='faster_rcnn_X_101_32x8d_FPN_3x',
# model='faster_rcnn_R_101_FPN_3x',
# model=model,
)
inferer_seg.cfg.INPUT.MIN_SIZE_TEST = size
inferer_seg.cfg.INPUT.MAX_SIZE_TEST = size
inferer_seg.cfg.MODEL.MASK_ON = True
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data7/weights/Keypoints_R_50_1x_aug_cm_seg_val_1/model_0009999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data7_0.1/weights/Keypoints_R_50_1x_aug_cm_seg_val_3/model_0009999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data7_0.1/weights/Keypoints_R_50_1x_aug_cm_seg_val_1/model_0007999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data8/weights/Keypoints_R_50_1x_aug_key_seg_val_1/model_0009999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data8/weights/Keypoints_R_50_1x_aug_key_seg_val_2/model_0004999.pth'
# inferer_key = jDetectron2KeypointInferer(
# weights_path=weights_path,
# # ref_coco_ann_path=f'/home/jitesh/3d/data/coco_data/hook_real1/json/hook.json',
# # categories_path=f'/home/jitesh/3d/data/categories/hook_infer.json',
# # categories_path=f'/home/jitesh/3d/data/categories/hook_7ckpt.json',
# categories_path=f'/home/jitesh/3d/data/categories/hook_7ckpt_pole.json',
# target_category='hook',
# model_name='keypoint_rcnn_R_50_FPN_1x',
# bbox_threshold=bbox_thresh,
# kpt_threshold=kpt_thresh,
# key_box='hook',
# )
# k_size = 1024
# inferer_key.cfg.INPUT.MIN_SIZE_TEST = k_size
# inferer_key.cfg.INPUT.MAX_SIZE_TEST = k_size
possible_modes = ['save', 'preview']
mode = 'save'
check_value(mode, valid_value_list=possible_modes)
# make_dir_if_not_exists(infer_dump_dir)
img_extensions = ['jpg', 'JPG', 'png', 'PNG']
img_pathlist = get_all_files_in_extension_list(
dir_path=f'{path}', extension_list=img_extensions)
img_pathlist.sort()
confirm_folder(infer_dump_dir, mode)
# confirm_folder(f'{infer_dump_dir}/good_seg', mode)
# confirm_folder(f'{infer_dump_dir}/good_cropped', mode)
# confirm_folder(f'{infer_dump_dir}/good', mode)
# confirm_folder(f'{infer_dump_dir}/G(>4D) P(>4D)', mode)
# confirm_folder(f'{infer_dump_dir}/G(>4D) P(<4D)', mode)
# confirm_folder(f'{infer_dump_dir}/G(<4D) P(>4D)', mode)
# confirm_folder(f'{infer_dump_dir}/G(<4D) P(<4D)', mode)
# confirm_folder(f'{infer_dump_dir}/bad', mode)
confirm_folder(f'{infer_dump_dir}/infer_key_seg', mode)
count = 0
start = datetime.now()
df = pd.DataFrame(data=[],
columns=[
'gt_d',
'pred_d',
'gt_ab',
'pred_ab',
'gt_ratio',
'pred_ratio',
'gt_ratio>4',
'pred_ratio>4',
'correct_above4d_ratio',
'incorrect_above4d_ratio',
'correct_below4d_ratio',
'incorrect_below4d_ratio',
])
# 'image_path'])
for i, img_path in enumerate(tqdm(
img_pathlist,
desc='Writing images',
)):
img_filename = get_filename(img_path)
# if not '201005_70_縮小革命PB020261.jpg' in img_path:
# continue
# if i > 19:
# continue
printj.purple(img_path)
img = cv2.imread(img_path)
result = img
# print(f'shape {img.shape}')
# cv2.imshow('i', img)
# cv2.waitKey(100000)
# continue
score_list, pred_class_list, bbox_list, pred_masks_list, pred_keypoints_list, vis_keypoints_list, kpt_confidences_list = inferer_seg.predict(
img=img)
# printj.blue(pred_masks_list)
max_hook_score = -1
max_pole_score = -1
diameter = -1
len_ab = -1
found_hook = False
found_pole = False
for score, pred_class, bbox, mask, keypoints, vis_keypoints, kpt_confidences in zip(
score_list, pred_class_list, bbox_list, pred_masks_list,
pred_keypoints_list, vis_keypoints_list, kpt_confidences_list):
if pred_class == 'pole':
found_pole = True
if max_pole_score < score:
# if True:
max_pole_score = score
diameter = compute_diameter(mask)
# result = draw_bool_mask(img=result, mask=mask, color=[
# 0, 255, 255],
# transparent=True
# )
pole_bbox_text = f'pole {str(round(score, 2))}'
pole_bbox = bbox
pole_mask = mask
# result = draw_bbox(img=result, bbox=bbox,
# text=pole_bbox_text, label_only=not show_bbox_border, label_orientation='bottom')
printj.blue(f'diameter={diameter}')
if pred_class == 'hook':
# printj.green.bold_on_yellow(score)
found_hook = True
if max_hook_score < score:
# if True:
max_hook_score = score
hook_bbox = BBox.buffer(bbox)
hook_score = round(score, 2)
hook_mask = mask
hook_keypoints = keypoints
hook_vis_keypoints = vis_keypoints
hook_kpt_confidences = kpt_confidences
# xmin, ymin, xmax, ymax = bbox.to_int().to_list()
# _xmin, _ymin, _xmax, _ymax = _bbox.to_int().to_list()
# width = _xmax-_xmin
# height = _ymax-_ymin
# scale = 0.2
# xmin = max(int(_xmin - width*scale), 0)
# xmax = min(int(_xmax + width*scale), img.shape[1])
# ymin = max(int(_ymin - height*scale), 0)
# ymax = min(int(_ymax + height*scale), img.shape[0])
# printj.red(score)
# printj.red(bbox)
# return
# img = draw_bbox(img=img, bbox=_bbox, color=[
# 0, 255, 255], thickness=2, text=f"{pred_class} {round(score, 3)}",
# label_orientation='top')
# img = draw_bbox(img=img, bbox=_bbox, color=[
# 0, 255, 255], thickness=2, text=f"{pred_class} {round(score, 3)}",
# label_orientation='bottom')
# result = draw_bool_mask(img=result, mask=mask, color=[
# 255, 255, 0],
# transparent=True
# )
# result = result
# bbox_text = str(round(score, 4))
# result = draw_bbox(img=result, bbox=bbox,
# text=bbox_text, label_only=not show_bbox_border)
bbox_label_mode = 'euler'
# result = draw_keypoints(
# img=result, keypoints=vis_keypoints, radius=2, color=[0, 0, 255],
# # keypoint_labels=kpt_labels, show_keypoints_labels=True, label_thickness=1,
# # ignore_kpt_idx=conf_idx_list
# )
kpt_labels = [
"kpt-a", "kpt-b", "kpt-cb", "kpt-c", "kpt-cd", "kpt-d",
"kpt-e"
]
kpt_skeleton = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5],
[5, 6]]
conf_idx_list = np.argwhere(
np.array(kpt_confidences) > conf_thresh).reshape(-1)
not_conf_idx_list = np.argwhere(
np.array(kpt_confidences) <= conf_thresh).reshape(
-1).astype(int)
conf_keypoints, conf_kpt_labels = np.array(vis_keypoints)[
conf_idx_list], np.array(kpt_labels)[conf_idx_list]
not_conf_keypoints, not_conf_kpt_labels = np.array(
vis_keypoints)[not_conf_idx_list], np.array(
kpt_labels)[not_conf_idx_list]
cleaned_keypoints = np.array(vis_keypoints.copy()).astype(
np.float32)
# result = draw_bool_mask(img=result, mask=mask, color=[
# 255, 255, 0],
# transparent=True
# )
# result, len_ab = draw_inference_on_hook2(img=result, cleaned_keypoints=cleaned_keypoints, kpt_labels=kpt_labels, kpt_skeleton=kpt_skeleton,
# score=score, bbox=_bbox, vis_keypoints=vis_keypoints, kpt_confidences=kpt_confidences, conf_idx_list=conf_idx_list, not_conf_idx_list=not_conf_idx_list,
# conf_keypoints=conf_keypoints, conf_kpt_labels=conf_kpt_labels, not_conf_keypoints=not_conf_keypoints, not_conf_kpt_labels=not_conf_kpt_labels,
# conf_thresh=conf_thresh, show_bbox_border=show_bbox_border, bbox_label_mode=bbox_label_mode, index_offset=0, diameter=diameter)
# result=result
# printj.green(_bbox)
# printj.green(_bbox.to_int())
# printj.green(_bbox.to_int().to_list())
printj.green.on_white(max_hook_score)
if found_pole:
result = draw_bool_mask(img=result,
mask=pole_mask,
color=[0, 255, 255],
transparent=True)
result = draw_bbox(img=result,
bbox=pole_bbox,
text=pole_bbox_text,
label_only=not show_bbox_border,
label_orientation='top')
result = draw_bbox(img=result,
bbox=pole_bbox,
text=pole_bbox_text,
label_only=not show_bbox_border,
label_orientation='bottom')
if found_hook:
result = draw_bool_mask(img=result,
mask=hook_mask,
color=[255, 255, 0],
transparent=True)
result, len_ab = draw_inference_on_hook2(
img=result,
cleaned_keypoints=cleaned_keypoints,
kpt_labels=kpt_labels,
kpt_skeleton=kpt_skeleton,
score=hook_score,
bbox=hook_bbox,
vis_keypoints=hook_vis_keypoints,
kpt_confidences=hook_kpt_confidences,
conf_idx_list=conf_idx_list,
not_conf_idx_list=not_conf_idx_list,
conf_keypoints=conf_keypoints,
conf_kpt_labels=conf_kpt_labels,
not_conf_keypoints=not_conf_keypoints,
not_conf_kpt_labels=not_conf_kpt_labels,
conf_thresh=conf_thresh,
show_bbox_border=show_bbox_border,
bbox_label_mode=bbox_label_mode,
index_offset=0,
diameter=diameter)
printj.purple(len_ab)
if len_ab == 0:
printj.green(keypoints)
result = draw_info_box(result, len_ab, diameter)
# img: np.ndarray, cleaned_keypoints, kpt_labels: List[str], kpt_skeleton: List[list],
# score: float, bbox: BBox, vis_keypoints: list, kpt_confidences: list, conf_idx_list: list, not_conf_idx_list: list,
# conf_keypoints, conf_kpt_labels, not_conf_keypoints, not_conf_kpt_labels,
# conf_thresh: float = 0.3, show_bbox_border: bool = False, bbox_label_mode: str = 'euler', index_offset: int = 0, diameter=1
# cv2.imshow('i', result)
# # cv2.imwrite('i', result)
# cv2.waitKey(10000)
# quit_flag = cv_simple_image_viewer(img=result, preview_width=1000)
# if quit_flag:
# break
# cv2.imwrite(f"{infer_dump_dir}/good_seg/{img_filename}", result)
cv2.imwrite(f"{infer_dump_dir}/infer_key_seg/{img_filename}", result)
def move(self,
dst_dataroot: str,
include_depth: bool = True,
include_RT: bool = False,
camera_path: str = None,
fps_path: str = None,
preserve_filename: bool = False,
use_softlink: bool = False,
ask_permission_on_delete: bool = True,
show_pbar: bool = True):
make_dir_if_not_exists(dst_dataroot)
delete_all_files_in_dir(dst_dataroot,
ask_permission=ask_permission_on_delete,
verbose=False)
processed_image_id_list = []
pbar = tqdm(total=len(self.annotations),
unit='annotation(s)',
leave=True) if show_pbar else None
if pbar is not None:
pbar.set_description('Moving Linemod Dataset Data')
for linemod_ann in self.annotations:
if not dir_exists(linemod_ann.data_root):
raise FileNotFoundError(
f"Couldn't find data_root at {linemod_ann.data_root}")
# Images
linemod_image = self.images.get(id=linemod_ann.image_id)[0]
if linemod_image.id not in processed_image_id_list:
img_path = f'{linemod_ann.data_root}/{get_filename(linemod_image.file_name)}'
if not file_exists(img_path):
raise FileNotFoundError(
f"Couldn't find image at {img_path}")
if preserve_filename:
dst_img_path = f'{dst_dataroot}/{get_filename(linemod_image.file_name)}'
if file_exists(dst_img_path):
raise FileExistsError(f"""
Image already exists at {dst_img_path}
Hint: Use preserve_filename=False to bypass this error.
""")
else:
dst_filename = f'{linemod_image.id}.{get_extension_from_filename(linemod_image.file_name)}'
linemod_image.file_name = dst_filename
dst_img_path = f'{dst_dataroot}/{dst_filename}'
if not use_softlink:
copy_file(src_path=img_path,
dest_path=dst_img_path,
silent=True)
else:
create_softlink(src_path=rel_to_abs_path(img_path),
dst_path=rel_to_abs_path(dst_img_path))
processed_image_id_list.append(linemod_image.id)
# Masks
if not file_exists(linemod_ann.mask_path):
raise FileNotFoundError(
f"Couldn't find mask at {linemod_ann.mask_path}")
mask_path = linemod_ann.mask_path
if preserve_filename:
dst_mask_path = f'{dst_dataroot}/{get_filename(linemod_ann.mask_path)}'
if file_exists(dst_mask_path):
raise FileExistsError(f"""
Mask already exists at {dst_mask_path}
Hint: Use preserve_filename=False to bypass this error.
""")
else:
mask_filename = get_filename(linemod_ann.mask_path)
dst_filename = f'{linemod_ann.id}_mask.{get_extension_from_filename(mask_filename)}'
dst_mask_path = f'{dst_dataroot}/{dst_filename}'
linemod_ann.mask_path = dst_mask_path
if not use_softlink:
copy_file(src_path=mask_path,
dest_path=dst_mask_path,
silent=True)
else:
create_softlink(src_path=rel_to_abs_path(mask_path),
dst_path=rel_to_abs_path(dst_mask_path))
# Depth
if include_depth and linemod_ann.depth_path is not None:
if not file_exists(linemod_ann.depth_path):
raise FileNotFoundError(
f"Couldn't find depth at {linemod_ann.depth_path}")
depth_path = linemod_ann.depth_path
if preserve_filename:
dst_depth_path = f'{dst_dataroot}/{get_filename(linemod_ann.depth_path)}'
if file_exists(dst_depth_path):
raise FileExistsError(f"""
Depth already exists at {dst_depth_path}
Hint: Use preserve_filename=False to bypass this error.
""")
else:
depth_filename = get_filename(linemod_ann.depth_path)
dst_filename = f'{linemod_ann.id}_depth.{get_extension_from_filename(depth_filename)}'
dst_depth_path = f'{dst_dataroot}/{dst_filename}'
linemod_ann.depth_path = dst_depth_path
if not use_softlink:
copy_file(src_path=depth_path,
dest_path=dst_depth_path,
silent=True)
else:
create_softlink(src_path=rel_to_abs_path(depth_path),
dst_path=rel_to_abs_path(dst_depth_path))
# RT pickle files
if include_RT:
rootname = get_rootname_from_path(mask_path)
if rootname.endswith('_mask'):
rootname = rootname.replace('_mask', '')
rt_filename = f'{rootname}_RT.pkl'
rt_path = f'{linemod_ann.data_root}/{rt_filename}'
if not file_exists(rt_path):
raise FileNotFoundError(
f"Couldn't find RT pickle file at {rt_path}")
if preserve_filename:
dst_rt_path = f'{dst_dataroot}/{rt_filename}'
if file_exists(dst_depth_path):
raise FileExistsError(f"""
RT pickle file already exists at {dst_rt_path}
Hint: Use preserve_filename=False to bypass this error.
""")
else:
dst_rt_filename = f'{linemod_ann.id}_RT.pkl'
dst_rt_path = f'{dst_dataroot}/{dst_rt_filename}'
if not use_softlink:
copy_file(src_path=rt_path,
dest_path=dst_rt_path,
silent=True)
else:
create_softlink(src_path=rel_to_abs_path(rt_path),
dst_path=rel_to_abs_path(dst_rt_path))
if pbar is not None:
pbar.update()
# Camera setting
if camera_path is not None:
if not file_exists(camera_path):
raise FileNotFoundError(
f"Couldn't find camera settings at {camera_path}")
dst_camera_path = f'{dst_dataroot}/{get_filename(camera_path)}'
if file_exists(dst_camera_path):
raise FileExistsError(
f'Camera settings already saved at {dst_camera_path}')
if not use_softlink:
copy_file(src_path=camera_path,
dest_path=dst_camera_path,
silent=True)
else:
create_softlink(src_path=rel_to_abs_path(camera_path),
dst_path=rel_to_abs_path(dst_camera_path))
# FPS setting
if fps_path is not None:
if not file_exists(fps_path):
raise FileNotFoundError(
f"Couldn't find FPS settings at {fps_path}")
dst_fps_path = f'{dst_dataroot}/{get_filename(fps_path)}'
if file_exists(dst_fps_path):
raise FileExistsError(
f'FPS settings already saved at {dst_fps_path}')
if not use_softlink:
copy_file(src_path=fps_path,
dest_path=dst_fps_path,
silent=True)
else:
create_softlink(src_path=rel_to_abs_path(fps_path),
dst_path=rel_to_abs_path(dst_fps_path))
if pbar is not None:
pbar.close()
def to_coco(self,
img_dir: str = None,
mask_dir: str = None,
coco_license: COCO_License = None,
check_paths: bool = True,
mask_lower_bgr: Tuple[int] = None,
mask_upper_bgr: Tuple[int] = (255, 255, 255),
show_pbar: bool = True) -> COCO_Dataset:
dataset = COCO_Dataset.new(
description='Dataset converted from Linemod to COCO format.')
dataset.licenses.append(
coco_license if coco_license is not None else COCO_License(
url=
'https://github.com/cm107/annotation_utils/blob/master/LICENSE',
name='MIT License',
id=len(dataset.licenses)))
coco_license0 = dataset.licenses[-1]
for linemod_image in self.images:
file_name = get_filename(linemod_image.file_name)
img_path = linemod_image.file_name if img_dir is None else f'{img_dir}/{file_name}'
if file_exists(img_path):
date_captured = get_ctime(img_path)
else:
if check_paths:
raise FileNotFoundError(
f"Couldn't find image at {img_path}")
date_captured = ''
coco_image = COCO_Image(license_id=coco_license0.id,
file_name=file_name,
coco_url=img_path,
width=linemod_image.width,
height=linemod_image.height,
date_captured=date_captured,
flickr_url=None,
id=linemod_image.id)
dataset.images.append(coco_image)
pbar = tqdm(total=len(self.annotations),
unit='annotation(s)') if show_pbar else None
if pbar is not None:
pbar.set_description('Converting Linemod to COCO')
for linemod_ann in self.annotations:
mask_filename = get_filename(linemod_ann.mask_path)
if mask_dir is not None:
mask_path = f'{mask_dir}/{mask_filename}'
if not file_exists(mask_path):
if check_paths:
raise FileNotFoundError(
f"Couldn't find mask at {mask_path}")
else:
seg = Segmentation()
else:
seg = Segmentation.from_mask_path(mask_path,
lower_bgr=mask_lower_bgr,
upper_bgr=mask_upper_bgr)
elif file_exists(linemod_ann.mask_path):
seg = Segmentation.from_mask_path(linemod_ann.mask_path,
lower_bgr=mask_lower_bgr,
upper_bgr=mask_upper_bgr)
elif img_dir is not None and file_exists(
f'{img_dir}/{mask_filename}'):
seg = Segmentation.from_mask_path(f'{img_dir}/{mask_filename}',
lower_bgr=mask_lower_bgr,
upper_bgr=mask_upper_bgr)
elif not check_paths:
seg = Segmentation()
else:
raise FileNotFoundError(f"""
Couldn't resolve mask_path for calculating segmentation.
Please either specify mask_dir or correct the mask paths
in your linemod dataset.
linemod_ann.id: {linemod_ann.id}
linemod_ann.mask_path: {linemod_ann.mask_path}
""")
if len(seg) > 0:
bbox = seg.to_bbox()
else:
xmin = int(
linemod_ann.corner_2d.to_numpy(demarcation=True)[:,
0].min())
xmax = int(
linemod_ann.corner_2d.to_numpy(demarcation=True)[:,
0].max())
ymin = int(
linemod_ann.corner_2d.to_numpy(demarcation=True)[:,
1].min())
ymax = int(
linemod_ann.corner_2d.to_numpy(demarcation=True)[:,
1].max())
bbox = BBox(xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax)
keypoints = Keypoint2D_List.from_point_list(linemod_ann.fps_2d,
visibility=2)
keypoints_3d = Keypoint3D_List.from_point_list(linemod_ann.fps_3d,
visibility=2)
num_keypoints = len(keypoints)
if linemod_ann.category_id not in [
cat.id for cat in dataset.categories
]:
linemod_cat = self.categories.get_obj_from_id(
linemod_ann.category_id)
cat_keypoints = list(
'abcdefghijklmnopqrstuvwxyz'.upper())[:num_keypoints]
cat_keypoints_idx_list = [
idx for idx in range(len(cat_keypoints))
]
cat_keypoints_idx_list_shift_left = cat_keypoints_idx_list[
1:] + cat_keypoints_idx_list[:1]
dataset.categories.append(
COCO_Category(
id=linemod_ann.category_id,
supercategory=linemod_cat.supercategory,
name=linemod_cat.name,
keypoints=cat_keypoints,
skeleton=[[start_idx, end_idx]
for start_idx, end_idx in zip(
cat_keypoints_idx_list,
cat_keypoints_idx_list_shift_left)]))
coco_ann = COCO_Annotation(
id=linemod_ann.id,
category_id=linemod_ann.category_id,
image_id=linemod_ann.image_id,
segmentation=seg,
bbox=bbox,
area=bbox.area(),
keypoints=keypoints,
num_keypoints=num_keypoints,
iscrowd=0,
keypoints_3d=keypoints_3d,
camera=COCO_Camera(f=[linemod_ann.K.fx, linemod_ann.K.fy],
c=[linemod_ann.K.cx, linemod_ann.K.cy],
T=[0, 0]))
dataset.annotations.append(coco_ann)
if pbar is not None:
pbar.update()
if pbar is not None:
pbar.close()
return dataset
