def from_dict(cls, ann_dict: dict, strict: bool = True) -> COCO_Annotation:
if strict:
check_required_keys(ann_dict,
required_keys=[
'segmentation', 'num_keypoints', 'area',
'iscrowd', 'keypoints', 'image_id', 'bbox',
'category_id', 'id'
])
return COCO_Annotation(
segmentation=Segmentation.from_list(ann_dict['segmentation'],
demarcation=False),
num_keypoints=ann_dict['num_keypoints'],
area=ann_dict['area'],
iscrowd=ann_dict['iscrowd'],
keypoints=Keypoint2D_List.from_list(ann_dict['keypoints'],
demarcation=False),
image_id=ann_dict['image_id'],
bbox=BBox.from_list(ann_dict['bbox'], input_format='pminsize'),
category_id=ann_dict['category_id'],
id=ann_dict['id'],
keypoints_3d=Keypoint3D_List.from_list(
ann_dict['keypoints_3d'], demarcation=False)
if 'keypoints_3d' in ann_dict else None,
camera=Camera.from_dict(ann_dict['camera_params'])
if 'camera_params' in ann_dict else None)
else:
check_required_keys(
ann_dict, required_keys=['id', 'category_id', 'image_id'])
return COCO_Annotation(
segmentation=Segmentation.from_list(ann_dict['segmentation'],
demarcation=False)
if 'segmentation' in ann_dict else None,
num_keypoints=ann_dict['num_keypoints']
if 'num_keypoints' in ann_dict else None,
area=ann_dict['area'] if 'area' in ann_dict else None,
iscrowd=ann_dict['iscrowd'] if 'iscrowd' in ann_dict else None,
keypoints=Keypoint2D_List.from_list(ann_dict['keypoints'],
demarcation=False)
if 'keypoints' in ann_dict else None,
image_id=ann_dict['image_id'],
bbox=BBox.from_list(ann_dict['bbox'], input_format='pminsize')
if 'bbox' in ann_dict else None,
category_id=ann_dict['category_id'],
id=ann_dict['id'],
keypoints_3d=Keypoint3D_List.from_list(
ann_dict['keypoints_3d'], demarcation=False)
if 'keypoints_3d' in ann_dict else None,
camera=Camera.from_dict(ann_dict['camera_params'])
if 'camera_params' in ann_dict else None)
def from_dict(cls, ann_dict: dict) -> Detectron2_Annotation:
check_required_keys(
ann_dict,
required_keys=['iscrowd', 'bbox', 'category_id', 'bbox_mode'])
if ann_dict['bbox_mode'] == BoxMode.XYWH_ABS:
bbox = BBox.from_list(ann_dict['bbox'], input_format='pminsize')
elif ann_dict['bbox_mode'] == BoxMode.XYXY_ABS:
bbox = BBox.from_list(ann_dict['bbox'], input_format='pminpmax')
else:
raise NotImplementedError
if 'keypoints' in ann_dict:
keypoints = Keypoint2D_List.from_list(
value_list=ann_dict['keypoints'], demarcation=False)
else:
keypoints = Keypoint2D_List()
return Detectron2_Annotation(iscrowd=ann_dict['iscrowd'],
bbox=bbox,
keypoints=keypoints,
category_id=ann_dict['category_id'],
segmentation=Segmentation.from_list(
points_list=ann_dict['segmentation']
if 'segmentation' in ann_dict else [],
demarcation=False),
bbox_mode=ann_dict['bbox_mode'])
def get_instance_segmentation(self,
img,
lower_thres: list = (0, 0, 0),
upper_thresh: list = (255, 255, 255)):
color_mask = cv2.inRange(src=img,
lowerb=lower_thres,
upperb=upper_thresh)
color_contours, _ = cv2.findContours(color_mask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
seg = Segmentation.from_contour(contour_list=color_contours)
return seg
def run(
path: str,
key: str = 'measure',
show_image: bool = False,
show_video: bool = True,
):
dataset = COCO_Dataset.load_from_path(json_path=f'{path}/{key}-coco.json',
img_dir=path,
strict=False)
# for i, coco_ann in enumerate(dataset.annotations):
# if i % 2 == 0:
# coco_ann.segmentation = Segmentation()
# if i % 3 == 0:
# coco_ann.keypoints = Keypoint2D_List()
# coco_ann.num_keypoints = 0
for coco_ann in dataset.annotations:
coco_ann.segmentation = Segmentation()
coco_ann.keypoints = Keypoint2D_List()
coco_ann.num_keypoints = 0
coco_ann.keypoints_3d = None
coco_ann.camera = None
for coco_cat in dataset.categories:
coco_cat.keypoints = []
coco_cat.skeleton = []
dataset.save_to_path('non_strict_dataset.json',
overwrite=True,
strict=False)
dataset0 = COCO_Dataset.load_from_path('non_strict_dataset.json',
strict=False)
dataset0.images.sort(attr_name='file_name')
if show_image:
dataset0.save_visualization(save_dir=f'{path}_{dt_string3}_visualize',
show_preview=True,
kpt_idx_offset=-1,
overwrite=True,
show_details=True)
if show_video:
dataset0.save_video(
save_path=f'{path}_{dt_string3}.mp4',
# draw_order=['screw'],
show_details=True,
show_preview=True,
kpt_idx_offset=-1,
overwrite=True,
fps=5,
show_seg=True)
logger.green('Visualisation complete')
def generate_visualizations(self, do_sort: bool = False):
delete_dir_if_exists(self.visualization_dump_dir)
make_dir(self.visualization_dump_dir)
if do_sort:
self.parser.images.image_list.sort(
key=operator.attrgetter('file_name'))
for coco_image in self.parser.images.image_list:
coco_image = COCO_Image.buffer(coco_image)
img_path = f"{self.img_dir}/{coco_image.file_name}"
img = cv2.imread(img_path)
coco_annotation_list = self.parser.annotations.get_annotations_from_imgIds(
imgIds=[coco_image.id])
bbox_list = []
kpt_list_list = []
kpt_skeleton_list = []
kpt_label_list_list = []
seg_list = []
cat_name_list = []
for coco_annotation in coco_annotation_list:
coco_annotation = COCO_Annotation.buffer(coco_annotation)
bbox = coco_annotation.bounding_box
keypoint_list = coco_annotation.get_keypoint_list()
kpt_list = [[kpt.x, kpt.y] for kpt in keypoint_list]
seg = Segmentation.from_list(
points_list=coco_annotation.segmentation)
bbox_list.append(bbox)
kpt_list_list.append(kpt_list)
seg_list.append(seg)
coco_category = self.parser.categories.get_category_from_id(
id=coco_annotation.category_id)
cat_name_list.append(coco_category.name)
kpt_skeleton_list.append(coco_category.skeleton)
kpt_label_list_list.append(coco_category.keypoints)
self.save(img=img,
bbox_list=bbox_list,
kpt_list_list=kpt_list_list,
kpt_skeleton_list=kpt_skeleton_list,
kpt_label_list_list=kpt_label_list_list,
seg_list=seg_list,
cat_name_list=cat_name_list,
file_name=coco_image.file_name)
if self.viz_limit is not None and self.viz_count >= self.viz_limit:
break
def __init__(
self,
id: int,
category_id: int,
image_id: int, # Standard Required
segmentation: Segmentation = None, # Standard Optional
bbox: BBox = None,
area: float = None,
keypoints: Keypoint2D_List = None,
num_keypoints: int = None,
iscrowd: int = 0,
keypoints_3d: Keypoint3D_List = None,
camera: Camera = None # Custom Optional
):
super().__init__(id=id)
# Standard Required
self.category_id = category_id
self.image_id = image_id
# Standard Optional
self.segmentation = segmentation if segmentation is not None else Segmentation(
[])
if bbox is not None:
self.bbox = bbox
else:
if len(self.segmentation) > 0:
self.bbox = self.segmentation.to_bbox()
else:
logger.error(
f'A COCO_Annotation needs to be given either a bbox or a non-empty segmentation at the very least to make a valid annotation.'
)
logger.error(
f'id: {id}, category_id: {category_id}, image_id: {image_id}'
)
raise Exception
self.area = area
self.keypoints = keypoints if keypoints is not None else Keypoint2D_List(
[])
self.num_keypoints = num_keypoints if num_keypoints is not None else len(
self.keypoints)
self.iscrowd = iscrowd
# Custom Optional
self.keypoints_3d = keypoints_3d
self.camera = camera
def get_instance_segmentation(self,
img: np.ndarray,
target_bgr: List[int] = None,
interval: int = 1,
exclude_invalid_polygons: bool = True):
target_bgr = target_bgr if target_bgr is not None else self.get_color_from_id(
)
check_list_length(target_bgr, correct_length=3)
lower_bgr = [
val - interval if val - interval >= 0 else 0 for val in target_bgr
]
upper_bgr = [
val + interval if val + interval <= 255 else 255
for val in target_bgr
]
color_mask = cv2.inRange(src=img,
lowerb=tuple(lower_bgr),
upperb=tuple(upper_bgr))
color_contours, _ = cv2.findContours(color_mask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
seg = Segmentation.from_contour(
contour_list=color_contours,
exclude_invalid_polygons=exclude_invalid_polygons)
return seg
measure_dir = 'measure_coco/measure'
whole_number_dir = 'measure_coco/whole_number'
digit_dir = 'measure_coco/digit'
json_output_filename = 'output.json'
measure_dataset, whole_number_dataset, digit_dataset = dataset.split_measure_dataset(
measure_dir=measure_dir,
whole_number_dir=whole_number_dir,
digit_dir=digit_dir,
allow_no_measures=True,
allow_missing_parts=True)
if False: # Change to True if you want to remove all segmentation from the measure dataset.
from common_utils.common_types.segmentation import Segmentation
for coco_ann in measure_dataset.annotations:
coco_ann.seg = Segmentation()
measure_dataset.display_preview(show_details=True,
window_name='Measure Dataset Preview')
logger.info(f'Saving Measure Dataset')
measure_dataset.save_to_path(f'{measure_dir}/{json_output_filename}',
overwrite=True)
whole_number_dataset.display_preview(
show_details=True, window_name='Whole Number Dataset Preview')
logger.info(f'Saving Whole Number Dataset')
whole_number_dataset.save_to_path(f'{whole_number_dir}/{json_output_filename}',
overwrite=True)
if False: # For debugging 2-digit digit annotations
del_ann_id_list = []
for segmentation_list in item.segmentation:
if len(segmentation_list.to_list()) < 5:
logger.red(coco_image.coco_url)
logger.red(item.segmentation)
logger.red(f"segmentation has abnormal point count {segmentation_list}" )
continue
poly_shapely = segmentation_list.to_shapely()
poly_list = list(zip(*poly_shapely.exterior.coords.xy))
line_non_simple = shapely.geometry.LineString(poly_list)
mls = shapely.ops.unary_union(line_non_simple)
polygons = list(shapely.ops.polygonize(mls))
poly_shapely = polygons
segmentation.append(Segmentation.from_shapely(poly_shapely))
image, bbox, poly = handler(image=img, bounding_boxes=bbox, polygons=segmentation)
poly = [Segmentation(polygon_list=[pol]) for pol in poly]
print("total bbox after")
for polygons in poly:
image = draw_segmentation(img=image, segmentation=polygons, transparent=True)
print("finished drawing handler")
cv2.imshow("Window", image)
cv2.waitKey(4000)
cv2.imwrite(os.path.join(save_img_path,str(i)+'.jpg'),image)
i+=1
from logger import logger
from common_utils.common_types.segmentation import Segmentation, Polygon
import numpy as np
contours = [
np.array([[[0, 0]], [[1, 0]], [[1, 2]], [[0, 1]]]),
np.array([[[0, 0]], [[1, 0]], [[1, 2]], [[0, 1]]]),
np.array([[[1, 0]], [[2, 0]], [[2, 2]], [[1, 1]]]),
np.array([[[1, 1]], [[2, 1]], [[2, 3]], [[2, 10]], [[1, 2]]]),
np.array([[[50, 100]], [[100, 200]]])
]
# seg = Segmentation.from_contour(contour_list=contours)
seg = Segmentation.from_contour(contour_list=contours,
exclude_invalid_polygons=True)
logger.purple(seg)
def _process_organized_part(self, organized_part: dict,
frame_img: np.ndarray,
whole_number_coco_image: COCO_Image,
measure_ann: COCO_Annotation,
whole_number_count: int, digit_dir: str):
whole_number_cat = self.whole_number_dataset.categories.get_unique_category_from_name(
'whole_number')
whole_number_seg = None
whole_number_abs_bbox = None
seg_len_list = [
len(part_ann.segmentation) for part_ann in organized_part['anns']
]
if all([seg_len > 0 for seg_len in seg_len_list]):
target = 'seg'
elif all([seg_len == 0 for seg_len in seg_len_list]):
target = 'bbox'
else:
logger.warning(
f"There are empty segmentations among organized_part['anns'].")
logger.warning(
f"Cannot determine whether to merge segmentations or bounding boxes."
)
logger.warning(
f"[len(part_ann.segmentation) for part_ann in organized_part['anns']]: {[len(part_ann.segmentation) for part_ann in organized_part['anns']]}"
)
logger.warning(
f"organized_part['whole_names']: {organized_part['whole_names']}"
)
logger.warning(
f"organized_part['part_names']: {organized_part['part_names']}"
)
logger.warning(f"Assuming target = 'bbox'")
target = 'bbox'
if target == 'seg':
for part_ann in organized_part['anns']:
whole_number_seg = whole_number_seg + part_ann.segmentation if whole_number_seg is not None else part_ann.segmentation
whole_number_abs_bbox = whole_number_seg.to_bbox().to_int()
elif target == 'bbox':
whole_number_seg = Segmentation()
for part_ann in organized_part['anns']:
whole_number_abs_bbox = whole_number_abs_bbox + part_ann.bbox if whole_number_abs_bbox is not None else part_ann.bbox
whole_number_abs_bbox = whole_number_abs_bbox.to_int()
else:
raise Exception
whole_number_bbox_region = frame_img[
whole_number_abs_bbox.ymin:whole_number_abs_bbox.ymax,
whole_number_abs_bbox.xmin:whole_number_abs_bbox.xmax, :]
whole_number_img_save_path = f'{digit_dir}/{get_rootname_from_filename(whole_number_coco_image.file_name)}_{whole_number_count}.{get_extension_from_filename(whole_number_coco_image.file_name)}'
self._save_image(img=whole_number_bbox_region,
save_path=whole_number_img_save_path)
digit_coco_image = COCO_Image.from_img_path(
img_path=whole_number_img_save_path,
license_id=0,
image_id=len(self.digit_dataset.images))
self.digit_dataset.images.append(digit_coco_image)
measure_bbox_ref_point = Point2D(x=measure_ann.bbox.xmin,
y=measure_ann.bbox.ymin)
if target == 'seg':
whole_number_seg = whole_number_seg - measure_bbox_ref_point # relative to measure bbox
whole_number_bbox = whole_number_seg.to_bbox(
) # relative to measure bbox
elif target == 'bbox':
whole_number_seg = Segmentation() # Empty
whole_number_bbox = whole_number_abs_bbox - measure_bbox_ref_point # relative to measure bbox
else:
raise Exception
whole_number_coco_ann = COCO_Annotation(
id=len(self.whole_number_dataset.annotations),
category_id=whole_number_cat.id,
image_id=whole_number_coco_image.id,
segmentation=whole_number_seg,
bbox=whole_number_bbox,
area=whole_number_bbox.area())
self.whole_number_dataset.annotations.append(whole_number_coco_ann)
for part_ann in organized_part['anns']:
part_ann = COCO_Annotation.buffer(part_ann)
coco_cat = self.categories.get_obj_from_id(part_ann.category_id)
digit_cat = self.digit_dataset.categories.get_unique_category_from_name(
coco_cat.name.split('part')[-1])
whole_number_bbox_ref_point = measure_bbox_ref_point + Point2D(
x=whole_number_bbox.xmin, y=whole_number_bbox.ymin)
if target == 'seg':
digit_seg = part_ann.segmentation - whole_number_bbox_ref_point
digit_bbox = digit_seg.to_bbox()
elif target == 'bbox':
digit_seg = Segmentation()
digit_bbox = part_ann.bbox - whole_number_bbox_ref_point
else:
raise Exception
digit_coco_ann = COCO_Annotation(id=len(
self.digit_dataset.annotations),
category_id=digit_cat.id,
image_id=digit_coco_image.id,
segmentation=digit_seg,
bbox=digit_bbox,
area=digit_bbox.area())
self.digit_dataset.annotations.append(digit_coco_ann)
seg = Segmentation(
[
Polygon.from_list(
points=[
[0, 0], [1, 0], [1, 2], [0, 1]
],
dimensionality=2,
demarcation=True
),
Polygon.from_list(
points=[
[0.5, 0], [1.5, 0], [1.5, 2], [0.5, 1]
],
dimensionality=2,
demarcation=True
),
Polygon.from_list(
points=[
[1, 0], [2, 0], [2, 2], [1, 1]
],
dimensionality=2,
demarcation=True
),
Polygon.from_list(
points=[
[1, 1], [2, 1], [2, 3], [2, 10], [1, 2]
],
dimensionality=2,
demarcation=True
)
]
)
Point3D_List([point * const_int for point in poly3d_0.to_point3d_list()]))
assert poly2d_0 / const_float == Polygon.from_point2d_list(
Point2D_List([point / const_float
for point in poly2d_0.to_point2d_list()]))
assert poly3d_0 / const_float == Polygon.from_point3d_list(
Point3D_List([point / const_float
for point in poly3d_0.to_point3d_list()]))
assert poly2d_0 / const_int == Polygon.from_point2d_list(
Point2D_List([point / const_int for point in poly2d_0.to_point2d_list()]))
assert poly3d_0 / const_int == Polygon.from_point3d_list(
Point3D_List([point / const_int for point in poly3d_0.to_point3d_list()]))
print('Polygon Test Passed')
seg0 = Segmentation([poly2d_0, poly2d_1])
seg1 = Segmentation([poly2d_2])
assert seg0 + seg1 == Segmentation(seg0.polygon_list + seg1.polygon_list)
assert seg0 + pt2d_0 == Segmentation([poly + pt2d_0 for poly in seg0])
assert seg0 + kpt2d_0 == Segmentation([poly + kpt2d_0 for poly in seg0])
assert seg0 + const_int == Segmentation([poly + const_int for poly in seg0])
assert seg0 + const_float == Segmentation(
[poly + const_float for poly in seg0])
assert seg0 - pt2d_0 == Segmentation([poly - pt2d_0 for poly in seg0])
assert seg0 - kpt2d_0 == Segmentation([poly - kpt2d_0 for poly in seg0])
assert seg0 - const_int == Segmentation([poly - const_int for poly in seg0])
assert seg0 - const_float == Segmentation(
[poly - const_float for poly in seg0])
def from_dict(cls, ann_dict: dict, strict: bool = True) -> COCO_Annotation:
if strict:
check_required_keys(ann_dict,
required_keys=[
'segmentation', 'num_keypoints', 'area',
'iscrowd', 'keypoints', 'image_id', 'bbox',
'category_id', 'id'
])
return COCO_Annotation(
segmentation=Segmentation.from_list(ann_dict['segmentation'],
demarcation=False),
num_keypoints=ann_dict['num_keypoints'],
area=ann_dict['area'],
iscrowd=ann_dict['iscrowd'],
keypoints=Keypoint2D_List.from_list(ann_dict['keypoints'],
demarcation=False),
image_id=ann_dict['image_id'],
bbox=BBox.from_list(ann_dict['bbox'], input_format='pminsize'),
category_id=ann_dict['category_id'],
id=ann_dict['id'],
keypoints_3d=Keypoint3D_List.from_list(
ann_dict['keypoints_3d'], demarcation=False)
if 'keypoints_3d' in ann_dict else None,
camera=Camera.from_dict(ann_dict['camera_params'])
if 'camera_params' in ann_dict else None)
else:
check_required_keys(
ann_dict, required_keys=['id', 'category_id', 'image_id'])
if 'segmentation' not in ann_dict:
seg = None
elif ann_dict['iscrowd'] == 1 and type(
ann_dict['segmentation']) == dict:
compressed_rle = mask.frPyObjects(
ann_dict['segmentation'],
ann_dict['segmentation']['size'][0],
ann_dict['segmentation']['size'][1])
seg_mask = mask.decode(compressed_rle)
contours, _ = cv2.findContours(seg_mask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
seg = Segmentation.from_contour(contour_list=contours)
else:
seg = Segmentation.from_list(ann_dict['segmentation'],
demarcation=False)
return COCO_Annotation(
segmentation=seg,
num_keypoints=ann_dict['num_keypoints']
if 'num_keypoints' in ann_dict else None,
area=ann_dict['area'] if 'area' in ann_dict else None,
iscrowd=ann_dict['iscrowd'] if 'iscrowd' in ann_dict else None,
keypoints=Keypoint2D_List.from_list(ann_dict['keypoints'],
demarcation=False)
if 'keypoints' in ann_dict else None,
image_id=ann_dict['image_id'],
bbox=BBox.from_list(ann_dict['bbox'], input_format='pminsize')
if 'bbox' in ann_dict else None,
category_id=ann_dict['category_id'],
id=ann_dict['id'],
keypoints_3d=Keypoint3D_List.from_list(
ann_dict['keypoints_3d'], demarcation=False)
if 'keypoints_3d' in ann_dict else None,
camera=Camera.from_dict(ann_dict['camera_params'])
if 'camera_params' in ann_dict else None)
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