def init_save_dir(self): make_dir_if_not_exists(self._save_dir) if self._clear: delete_all_files_in_dir(self._save_dir, ask_permission=False) else: self._existing_extensions = list(set([get_extension_from_path(path) for path in get_valid_image_paths(self._save_dir)])) self._first = False
def save_to_dir(self, json_save_dir: str, src_img_dir: str, dst_img_dir: str = None, overwrite: bool = False, show_pbar: bool = False): """Saves NDDS_Frame_Handler object to a directory path. Arguments: json_save_dir {str} -- [Path to directory where you want to save the NDDS annotation json files.] src_img_dir {str} -- [Path to directory where the original NDDS images are saved.] Keyword Arguments: dst_img_dir {str} -- [Path to directory where you want to copy the original NDDS images.] (default: {None}) overwrite {bool} -- [Whether or not you would like to overwrite existing files/directories.] (default: {False}) show_pbar {bool} -- [Whether or not you would like to show the progress bar.] (default: {False}) """ self._check_paths_valid(src_img_dir=src_img_dir) make_dir_if_not_exists(json_save_dir) delete_all_files_in_dir(json_save_dir, ask_permission=not overwrite) if dst_img_dir is not None: make_dir_if_not_exists(dst_img_dir) delete_all_files_in_dir(dst_img_dir, ask_permission=not overwrite) if show_pbar: pbar = tqdm(total=len(self), unit='ann(s)', leave=True) pbar.set_description(f'Saving {self.__class__.__name__}') for frame in self: save_path = f'{json_save_dir}/{get_rootname_from_path(frame.img_path)}.json' if dst_img_dir is not None: copy_file( src_path=f'{src_img_dir}/{get_filename(frame.img_path)}', dest_path=f'{dst_img_dir}/{get_filename(frame.img_path)}', silent=True) if frame.cs_img_path: copy_file( src_path= f'{src_img_dir}/{get_filename(frame.cs_img_path)}', dest_path= f'{dst_img_dir}/{get_filename(frame.cs_img_path)}', silent=True) if frame.depth_img_path: copy_file( src_path= f'{src_img_dir}/{get_filename(frame.depth_img_path)}', dest_path= f'{dst_img_dir}/{get_filename(frame.depth_img_path)}', silent=True) if frame.is_img_path: copy_file( src_path= f'{src_img_dir}/{get_filename(frame.is_img_path)}', dest_path= f'{dst_img_dir}/{get_filename(frame.is_img_path)}', silent=True) frame.ndds_ann.save_to_path(save_path=save_path) else: frame.ndds_ann.save_to_path(save_path=save_path) if show_pbar: pbar.update()
def setup_directories(self, verbose: bool=False): make_dir_if_not_exists(self.dest_dir) for target_dir in [self.train_dir, self.test_dir, self.val_dir]: if dir_exists(target_dir): delete_all_files_in_dir(dir_path=target_dir, ask_permission=False, verbose=verbose) else: make_dir(target_dir) img_dir = f"{target_dir}/{self.target_img_dir}" ann_dir = f"{target_dir}/{self.target_ann_dir}" make_dir(img_dir) make_dir(ann_dir)
def save_to_dir(self, json_save_dir: str, src_img_dir: str, overwrite: bool = False, dst_img_dir: str = None, show_pbar: bool = True): self._check_paths_valid(src_img_dir=src_img_dir) make_dir_if_not_exists(json_save_dir) delete_all_files_in_dir(json_save_dir, ask_permission=not overwrite) if dst_img_dir is not None: make_dir_if_not_exists(dst_img_dir) delete_all_files_in_dir(dst_img_dir, ask_permission=not overwrite) if show_pbar: pbar = tqdm(total=len(self), unit='ann(s)', leave=True) pbar.set_description(f'Saving {self.__class__.__name__}') for frame in self: save_path = f'{json_save_dir}/{get_rootname_from_path(frame.img_path)}.json' if dst_img_dir is not None: copy_file( src_path=f'{src_img_dir}/{get_filename(frame.img_path)}', dest_path=f'{dst_img_dir}/{get_filename(frame.img_path)}', silent=True) if frame.cs_img_path: copy_file( src_path= f'{src_img_dir}/{get_filename(frame.cs_img_path)}', dest_path= f'{dst_img_dir}/{get_filename(frame.cs_img_path)}', silent=True) if frame.depth_img_path: copy_file( src_path= f'{src_img_dir}/{get_filename(frame.depth_img_path)}', dest_path= f'{dst_img_dir}/{get_filename(frame.depth_img_path)}', silent=True) if frame.is_img_path: copy_file( src_path= f'{src_img_dir}/{get_filename(frame.is_img_path)}', dest_path= f'{dst_img_dir}/{get_filename(frame.is_img_path)}', silent=True) frame.ndds_ann.save_to_path(save_path=save_path) else: frame.ndds_ann.save_to_path(save_path=save_path) if show_pbar: pbar.update()
def save_to_dir(self, json_save_dir: str, src_img_dir: str, overwrite: bool=False, dst_img_dir: str=None): self._check_paths_valid(src_img_dir=src_img_dir) make_dir_if_not_exists(json_save_dir) delete_all_files_in_dir(json_save_dir, ask_permission=not overwrite) if dst_img_dir is not None: make_dir_if_not_exists(dst_img_dir) delete_all_files_in_dir(dst_img_dir, ask_permission=not overwrite) for ann in tqdm(self, total=len(self), unit='ann', leave=True): save_path = f'{json_save_dir}/{get_rootname_from_path(ann.img_path)}.json' src_img_path = f'{src_img_dir}/{get_filename(ann.img_path)}' if dst_img_dir is not None: dst_img_path = f'{dst_img_dir}/{get_filename(ann.img_path)}' copy_file(src_path=src_img_path, dest_path=dst_img_path, silent=True) ann.save_to_path(save_path=save_path, img_path=dst_img_path) else: ann.save_to_path(save_path=save_path, img_path=src_img_path)
def save_to_dir(self, json_save_dir: str, src_img_dir: str, overwrite: bool=False, dst_img_dir: str=None, show_pbar: bool=True): self._check_paths_valid(src_img_dir=src_img_dir) make_dir_if_not_exists(json_save_dir) delete_all_files_in_dir(json_save_dir, ask_permission=not overwrite) if dst_img_dir is not None: make_dir_if_not_exists(dst_img_dir) delete_all_files_in_dir(dst_img_dir, ask_permission=not overwrite) pbar = tqdm(total=len(self), unit='annotation(s)', leave=True) if show_pbar else None if pbar is not None: pbar.set_description('Writing Labelme Annotations') for ann in self: save_path = f'{json_save_dir}/{get_rootname_from_path(ann.img_path)}.json' src_img_path = f'{src_img_dir}/{get_filename(ann.img_path)}' if dst_img_dir is not None: dst_img_path = f'{dst_img_dir}/{get_filename(ann.img_path)}' copy_file(src_path=src_img_path, dest_path=dst_img_path, silent=True) ann.save_to_path(save_path=save_path, img_path=dst_img_path) else: ann.save_to_path(save_path=save_path, img_path=src_img_path) if pbar is not None: pbar.update() if pbar is not None: pbar.close()
from logger import logger from annotation_utils.ndds.structs import NDDS_Dataset from annotation_utils.coco.structs import COCO_Dataset, COCO_Category_Handler from common_utils.file_utils import make_dir_if_not_exists, delete_all_files_in_dir src_dir = '/home/clayton/workspace/prj/data_keep/data/ndds/20200924_yagura5_bright_1-colored' dst_dir = 'nihonbashi_debug' make_dir_if_not_exists(dst_dir) delete_all_files_in_dir(dst_dir, ask_permission=False) hsr_categories = COCO_Category_Handler.load_from_path( '/home/clayton/workspace/prj/data_keep/data/toyota/dataset/config/categories/hsr_categories.json' ) # Load NDDS Dataset ndds_dataset = NDDS_Dataset.load_from_dir(json_dir=src_dir, show_pbar=True) # Fix NDDS Dataset naming so that it follows convention. (This is not necessary if the NDDS dataset already follows the naming convention.) for frame in ndds_dataset.frames: # Fix Naming Convention for ann_obj in frame.ndds_ann.objects: if ann_obj.class_name.lower() == 'nihonbashi': obj_type, obj_name = 'seg', 'hsr' instance_name = '0' ann_obj.class_name = f'{obj_type}_{obj_name}_{instance_name}' elif ann_obj.class_name.lower() in list('abcdefghijkl'): obj_type, obj_name = 'kpt', 'hsr' instance_name, contained_name = '0', ann_obj.class_name ann_obj.class_name = f'{obj_type}_{obj_name}_{instance_name}_{contained_name}' else: logger.error(f'Unknown ann_obj.class_name: {ann_obj.class_name}')
ignore_unspecified_categories=True, show_pbar=True, bbox_area_threshold=1, default_visibility_threshold=0.10, visibility_threshold_dict={'measure': 0.01}, allow_unfound_seg=True, class_merge_map={ 'mark_10th_place': 'seg_measure', 'marking_bottom': 'seg_measure', 'marking_top': 'seg_measure', 'hook': 'seg_measure' }) # Output Directories make_dir_if_not_exists('measure_coco') delete_all_files_in_dir('measure_coco') 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
from logger import logger from annotation_utils.ndds.structs import NDDS_Dataset from annotation_utils.coco.structs import COCO_Dataset, COCO_Category_Handler from common_utils.file_utils import make_dir_if_not_exists, delete_all_files_in_dir src_root_dir = '/home/doors/workspace/oobayashi/hsr_data_4k' targets = [ '20200924_yagura5_bright_2-colored', 'yagura5_dark_2', 'yagura5_bright_2', 'yagura5_dark_1', '20200924_yagura5_dark_2-colored', '20200924_yagura5_bright_1-colored', '20200924_yagura5_dark_1-colored', 'yagura5_bright_1' ] dst_root_dir = '/home/doors/workspace/prj/data_keep/data/toyota/dataset/sim/20200928' video_preview_dir = f'{dst_root_dir}/preview' make_dir_if_not_exists(video_preview_dir) delete_all_files_in_dir(video_preview_dir, ask_permission=False) hsr_categories = COCO_Category_Handler.load_from_path( '/home/doors/workspace/prj/data_keep/data/toyota/dataset/config/categories/hsr_categories.json' ) for target in targets: src_target_dir = f'{src_root_dir}/{target}' dst_target_dir = f'{dst_root_dir}/{target}' make_dir_if_not_exists(dst_target_dir) delete_all_files_in_dir(dst_target_dir, ask_permission=False) # Load NDDS Dataset ndds_dataset = NDDS_Dataset.load_from_dir(json_dir=src_target_dir, show_pbar=True)
from annotation_utils.coco.structs import COCO_Category_Handler, COCO_Category from annotation_utils.coco.structs import COCO_Dataset from common_utils.file_utils import make_dir_if_not_exists, delete_all_files_in_dir from typing import cast from annotation_utils.coco.structs import COCO_Image, COCO_Annotation from common_utils.path_utils import get_rootname_from_filename, get_extension_from_filename from common_utils.common_types.point import Point2D from tqdm import tqdm import cv2 from logger import logger from typing import List target_src_dir = '/home/clayton/workspace/prj/data_keep/data/ndds/bolt_markMap_2020.08.18-13.03.35' target_dst_dir = 'bolt_kpt' make_dir_if_not_exists(target_dst_dir) delete_all_files_in_dir(target_dst_dir) # Load NDDS Dataset logger.info('Loading NDDS Dataset') ndds_dataset = NDDS_Dataset.load_from_dir( json_dir=target_src_dir, show_pbar=True ) delete_idx_list = [] # Fix NDDS Dataset naming so that it follows convention. (This is not necessary if the NDDS dataset already follows the naming convention.) for i, frame in enumerate(ndds_dataset.frames): for ann_obj in frame.ndds_ann.objects: if ann_obj.class_name.startswith('bolt'): if ann_obj.visibility == 0 :
import cv2 from annotation_utils.ndds.structs import NDDS_Dataset from common_utils.file_utils import make_dir_if_not_exists, delete_all_files_in_dir from common_utils.cv_drawing_utils import draw_bbox from streamer.cv_viewer import cv_simple_image_viewer # Load NDDS Dataset ndds_dataset = NDDS_Dataset.load_from_dir( json_dir='/home/clayton/workspace/prj/data_keep/data/ndds/m1_200', show_pbar=True) vis_dump_dir = 'vis_dump' make_dir_if_not_exists(vis_dump_dir) delete_all_files_in_dir(vis_dump_dir) for frame in ndds_dataset.frames: img = cv2.imread(frame.img_path) for ndds_obj in frame.ndds_ann.objects: img = draw_bbox(img=img, bbox=ndds_obj.bounding_box) quit_flag = cv_simple_image_viewer(img=img, preview_width=1000) if quit_flag: break
def _prep_output_dir(self, measure_dir: str, whole_number_dir: str, digit_dir: str): for output_dir in [measure_dir, whole_number_dir, digit_dir]: make_dir_if_not_exists(output_dir) delete_all_files_in_dir(output_dir, ask_permission=False)
def prepare_datasets_from_dir( scenario_root_dir: str, dst_root_dir: str, annotation_filename: str = 'output.json', skip_existing: bool = False, val_target_proportion: float = 0.05, min_val_size: int = None, max_val_size: int = None, orig_config_save: str = 'orig.yaml', reorganized_config_save: str = 'dataset_config.yaml', show_pbar: bool = True): """ Parameters: scenario_root_dir - Path to the source root directory containing all of your scenario folders. [Required] dst_root_dir - Path to where you would like to save your prepared scenario datasets (split into train and val) [Required] annotation_filename - The filename of every annotation file under scenario_root_dir [Default: 'output.json'] skip_existing - If you terminated dataset preparation midway, you can skip the scenarios that were already made using skip_existing=True. [Default: False] val_target_proportion - The proportion of your scenario that you would like to allocate to validation. [Default: 0.05] min_val_size - The minimum number of images that you would like to use for validation. [Default: None] max_val_size - The maximum number of images that you would like to use for validation. [Default: None] orig_config_save - Where you would like to save the dataset configuration representing your scenario_root_dir. [Default: 'orig.yaml] reorganized_config_save - Where you would like to save the dataset configuration representing your dst_root_dir. [Default: 'dataset_config.yaml'] show_pbar - Whether or not you would like to show a progress bar during preparation. [Default: True] Description: The datasets under each scenario directory will be combined and then split into a train + validation folder. The source root directory should have the following structure: scenario_root_dir scenario0 scenario0_part0 scenario0_part1 scenario0_part2 ... scenario1 scenario1_part0 scenario1_part1 scenario1_part2 ... scenario2 scenario2_part0 scenario2_part1 scenario2_part2 ... ... Note that there is no restriction on directory names, so the directory names should be anything. This method reads a fixed directory structure regardless of the directory names. Also note that it is necessary for the coco annotation file saved in each scenario part directory to have the same filename. If you need a more flexible approach for preparing your datasets, please define where your datasets are located in an excel sheet and use prepare_datasets_from_excel instead. The destination root directory will have the following structure: dst_root_dir scenario0 train val scenario1 train val scenario2 train val ... The dataset configuration file saved at reorganized_config_save will reflect the directory structure of dst_root_dir. The configuration file representing the directory structure of the scenario_root_dir is saved under orig_config_save. Note that orig_config_save and reorganized_config_save do not have to be inside of dst_root_dir. On the contrary, it is recommended to not save orig_config_save and reorganized_config_save inside of dst_root_dir. It is recommended that you change the path of orig_config_save and reorganized_config_save everytime you make an addition to your datasets. This is because you will likely want to keep track of the previous states of your dataset configuration, and you may also want to rollback to a previous configuration at any given time. """ make_dir_if_not_exists(dst_root_dir) if get_dir_contents_len(dst_root_dir) > 0 and not skip_existing: print( f'Directory {dst_root_dir} is not empty.\nAre you sure you want to delete the contents?' ) answer = input('yes/no') if answer.lower() in ['yes', 'y']: delete_all_files_in_dir(dst_root_dir) elif answer.lower() in ['no', 'n']: print(f'Terminating program.') sys.exit() else: raise ValueError(f'Invalid answer: {answer}') # Gather datasets from source root directory and combine. scenario_names = get_dirnames_in_dir(scenario_root_dir) scenario_datasets = cast(List[COCO_Dataset], []) orig_collection_handler = DatasetConfigCollectionHandler() pbar = tqdm(total=len(scenario_names), unit='scenario(s)') if show_pbar else None if pbar is not None: pbar.set_description('Gathering Scenarios') for scenario_name in scenario_names: orig_scenario_collection = DatasetConfigCollection(tag=scenario_name) src_scenario_dir = f'{scenario_root_dir}/{scenario_name}' part_names = get_dirnames_in_dir(src_scenario_dir) part_datasets = cast(List[COCO_Dataset], []) part_dataset_dirs = cast(List[str], []) for part_name in part_names: src_part_dir = f'{src_scenario_dir}/{part_name}' src_part_ann_path = f'{src_part_dir}/{annotation_filename}' part_dataset = COCO_Dataset.load_from_path( json_path=src_part_ann_path, img_dir=src_part_dir) part_datasets.append(part_dataset) part_dataset_dirs.append(src_part_dir) orig_scenario_part_config = DatasetConfig( img_dir=src_part_dir, ann_path=src_part_ann_path, ann_format='coco', tag=part_name) orig_scenario_collection.append(orig_scenario_part_config) scenario_dataset = COCO_Dataset.combine(dataset_list=part_datasets, img_dir_list=part_dataset_dirs, show_pbar=False) scenario_datasets.append(scenario_dataset) orig_collection_handler.append(orig_scenario_collection) if pbar is not None: pbar.update() orig_collection_handler.save_to_path(orig_config_save, overwrite=True) pbar.close() # Split combined scenario datasets into train and val and save them. train_collection = DatasetConfigCollection(tag='train') val_collection = DatasetConfigCollection(tag='val') pbar = tqdm(total=len(scenario_names)) if show_pbar else None if pbar is not None: pbar.set_description('Splitting Scenarios Into Train/Val') for i in range(len(scenario_names)): dst_scenario_dir = f'{dst_root_dir}/{scenario_names[i]}' if dir_exists(dst_scenario_dir): if skip_existing: if pbar is not None: pbar.update() continue else: raise FileExistsError( f'Directory already exists: {dst_scenario_dir}') else: make_dir_if_not_exists(dst_scenario_dir) orig_num_images = len(scenario_datasets[i].images) assert orig_num_images >= 2, f'{scenario_names[i]} has only {orig_num_images} images, and thus cannot be split into train and val.' num_val = int(len(scenario_datasets[i].images) * val_target_proportion) num_val = 1 if num_val == 0 else num_val num_val = min_val_size if min_val_size is not None and num_val < min_val_size else num_val num_val = max_val_size if max_val_size is not None and num_val > max_val_size else num_val num_train = orig_num_images - num_val train_dataset, val_dataset = scenario_datasets[i].split_into_parts( ratio=[num_train, num_val], shuffle=True) dst_train_dir = f'{dst_scenario_dir}/train' make_dir_if_not_exists(dst_train_dir) train_dataset.move_images(dst_img_dir=dst_train_dir, preserve_filenames=False, update_img_paths=True, show_pbar=False) train_ann_path = f'{dst_train_dir}/output.json' train_dataset.save_to_path(train_ann_path, overwrite=True) train_dataset_config = DatasetConfig(img_dir=dst_train_dir, ann_path=train_ann_path, ann_format='coco', tag=f'{scenario_names[i]}_train') train_collection.append(train_dataset_config) dst_val_dir = f'{dst_scenario_dir}/val' make_dir_if_not_exists(dst_val_dir) val_dataset.move_images(dst_img_dir=dst_val_dir, preserve_filenames=False, update_img_paths=True, show_pbar=False) val_ann_path = f'{dst_val_dir}/output.json' val_dataset.save_to_path(val_ann_path, overwrite=True) val_dataset_config = DatasetConfig(img_dir=dst_val_dir, ann_path=val_ann_path, ann_format='coco', tag=f'{scenario_names[i]}_val') val_collection.append(val_dataset_config) if pbar is not None: pbar.update() pbar.close() collection_handler = DatasetConfigCollectionHandler( [train_collection, val_collection]) collection_handler.save_to_path(reorganized_config_save, overwrite=True)
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 gen_infer_comparison(gt: BasicLoadableHandler, dt: BasicLoadableHandler, error: BasicLoadableHandler, model_names: List[str], test_names: List[str], collage_shape: (int, int), test_img_dir_map: Dict[str, str] = None, model_aliases: Dict[str, str] = None, test_aliases: Dict[str, str] = None, video_save: str = None, img_dump_dir: str = None, show_preview: bool = False, show_pbar: bool = True, draw_settings=None, draw_inference: bool = False, details_func=None, debug_verbose: bool = False): for handler in [gt, dt, error]: assert isinstance(handler, BasicLoadableHandler) for attr_key in ['frame', 'test_name']: assert hasattr(handler[0], attr_key) for handler in [dt, error]: assert hasattr(handler[0], 'model_name') model_names0 = list(set([datum.model_name for datum in dt ])) if model_names == 'all' else model_names test_names0 = list(set([datum.test_name for datum in gt ])) if test_names == 'all' else test_names for val_list in [model_names0, test_names0]: if val_list != 'all': assert isinstance(val_list, (tuple, list)) for val in val_list: assert isinstance(val, str) assert isinstance(collage_shape, (tuple, list)) for val in collage_shape: assert isinstance(val, int) assert len(collage_shape) == 2 assert len(model_names0) <= collage_shape[0] * collage_shape[1] if img_dump_dir is not None: make_dir_if_not_exists(img_dump_dir) delete_all_files_in_dir(img_dump_dir, ask_permission=False) if test_img_dir_map is None: test_img_dir_map0 = {test_name: test_name for test_name in test_names0} else: assert isinstance(test_img_dir_map, dict) for key, val in test_img_dir_map.items(): assert key in test_names0 assert isinstance(key, str) assert isinstance(val, str) test_img_dir_map0 = { test_name: (test_img_dir_map[test_name] if test_name in test_img_dir_map else test_name) \ for test_name in test_names0 } for test_name, img_dir in test_img_dir_map0.items(): if not dir_exists(img_dir): raise FileNotFoundError(f""" Couldn't find image directory {img_dir} for {test_name}. Please modify test_img_dir_map to match the image directory path for {test_name}. test_img_dir_map: {test_img_dir_map0} """) stream_writer = StreamWriter(show_preview=show_preview, video_save_path=video_save, dump_dir=img_dump_dir) total_images = len(gt.get(test_name=test_names0)) pbar = tqdm(total=total_images, unit='image(s)', leave=True) if show_pbar else None if pbar is not None: pbar.set_description('Generating Comparison') for test_name in test_names0: if img_dump_dir is not None: test_img_dump_dir = f'{img_dump_dir}/{test_name}' make_dir_if_not_exists(test_img_dump_dir) stream_writer.dump_writer._save_dir = test_img_dump_dir img_dir = test_img_dir_map0[test_name] gt_test_data = gt.get(test_name=test_name) gt_test_data.sort(attr_name='frame') dt_test_data = dt.get(test_name=test_name) dt_test_data.sort(attr_name='frame') error_test_data = error.get(test_name=test_name) error_test_data.sort(attr_name='frame') for gt_datum in gt_test_data: file_name = gt_datum.frame img_path = f'{img_dir}/{file_name}' if not file_exists(img_path): if debug_verbose: print(f""" Couldn't find image. Skipping. test_name: {test_name} img_path: {img_path} """) pbar.update() continue img = cv2.imread(img_path) dt_frame_data = dt_test_data.get(frame=gt_datum.frame) error_frame_data = error_test_data.get(frame=gt_datum.frame) img_buffer = cast(List[np.ndarray], []) for model_name in model_names0: dt_model_data = dt_frame_data.get(model_name=model_name) dt_model_datum = dt_model_data[0] if len( dt_model_data) > 0 else None error_model_data = error_frame_data.get(model_name=model_name) error_datum = error_model_data[0] if len( error_model_data) > 0 else None result = img.copy() if draw_inference or draw_settings is not None: if draw_settings is not None: if dt_model_datum is not None: result = dt_model_datum.draw( result, settings=draw_settings) else: if dt_model_datum is not None: result = dt_model_datum.draw(result) if test_aliases is not None and gt_datum.test_name in test_aliases: test_text = test_aliases[gt_datum.test_name] else: test_text = gt_datum.test_name if gt_datum is not None else None if model_aliases is not None and dt_model_datum.model_name in model_aliases: model_text = model_aliases[dt_model_datum.model_name] else: model_text = dt_model_datum.model_name if dt_model_datum is not None else None if details_func is not None: for key in ['gt', 'dt', 'error']: assert key in details_func.__annotations__, f'{details_func.__name__} must have a {key} parameter.' details_func_params = { 'img': result, 'gt': gt_datum, 'dt': dt_model_datum, 'error': error_datum } suggested_params = { 'test_text': test_text, 'model_text': model_text, 'frame_text': gt_datum.frame } for key, val in suggested_params.items(): if key in details_func.__annotations__: details_func_params[key] = val result = details_func(**details_func_params) else: row_text_list = [ f'Test: {test_text}', f'Model: {model_text}', f'Frame: {gt_datum.frame}' ] result_h, result_w = result.shape[:2] combined_row_height = len(row_text_list) * 0.04 * result_h result = draw_text_rows_at_point( img=result, row_text_list=row_text_list, x=result_w * 0.01, y=result_h * 0.01, combined_row_height=combined_row_height) img_buffer.append(result) collage_img = collage_from_img_buffer(img_buffer=img_buffer, collage_shape=collage_shape) stream_writer.step(img=collage_img, file_name=file_name) if pbar is not None: pbar.update() if pbar is not None: pbar.close()