def save(self, img: np.ndarray, bbox_list: list, kpt_list_list: list,
kpt_skeleton_list: list, kpt_label_list_list: list,
seg_list: list, cat_name_list: list, file_name: str):
result = img.copy()
save_path = f"{self.visualization_dump_dir}/{file_name}"
retry_count = 0
while file_exists(save_path):
rootname = get_rootname_from_filename(file_name)
extension = get_extension_from_filename(file_name)
save_path = f"{self.visualization_dump_dir}/{rootname}_{retry_count}.{extension}"
if retry_count == 9:
logger.error(f"Can't resolve save_path.")
raise Exception
retry_count += 1
for bbox, kpt_list, kpt_skeleton, kpt_label_list, seg, cat_name in \
zip(bbox_list, kpt_list_list, kpt_skeleton_list, kpt_label_list_list, seg_list, cat_name_list):
result = self._draw(img=result,
bbox=bbox,
kpt_list=kpt_list,
kpt_skeleton=kpt_skeleton,
kpt_label_list=kpt_label_list,
seg=seg,
cat_name=cat_name)
cv2.imwrite(filename=save_path, img=result)
logger.info(f"Wrote {save_path}")
self.viz_count += 1
def get_extensions(self) -> List[str]:
extension_list = []
for coco_image in self:
extension = get_extension_from_filename(coco_image.file_name)
if extension not in extension_list:
extension_list.append(extension)
return extension_list
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
)
crop_pbar = tqdm(total=len(bolt_roi_dataset.images), unit='image(s)')
crop_pbar.set_description('Cropping')
for roi_image in bolt_roi_dataset.images:
orig_img = cv2.imread(roi_image.coco_url)
roi_anns = bolt_roi_dataset.annotations.get_annotations_from_imgIds([roi_image.id])
mark_images = mark_dataset.images.get_images_from_file_name(roi_image.file_name)
assert len(mark_images) == 1
mark_image = mark_images[0]
mark_anns = mark_dataset.annotations.get_annotations_from_imgIds(mark_image.id)
#assert len(roi_anns) == len(mark_anns)
for i, roi_ann in enumerate(roi_anns):
roi_img = roi_ann.bbox.crop_from(orig_img)
img_rootname = get_rootname_from_filename(roi_image.file_name)
img_extension = get_extension_from_filename(roi_image.file_name)
save_filename = f'{img_rootname}_{i}.{img_extension}'
save_path = f'{marker_dst_dir}/{save_filename}'
cv2.imwrite(save_path, roi_img)
cropped_coco_image = COCO_Image.from_img_path(
img_path=save_path,
license_id=cropped_mark_dataset.licenses[0].id,
image_id=len(cropped_mark_dataset.images)
)
cropped_mark_dataset.images.append(cropped_coco_image)
mark_ann_found = False
mark_ann_list = cast(List[COCO_Annotation], [])
for i in list(range(len(mark_anns)))[::-1]:
if roi_ann.bbox.contains(mark_anns[i].bbox):
mark_ann_found = True