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