def save_to_path(self, save_path: str, overwrite: bool=False, strict: bool=True):
"""
Save this COCO_Dataset object to a json file in the standard COCO format.
save_path: Path of where you would like to save the dataset.
overwrite: If True, any existing file that exists at save_path will be overwritten.
"""
if file_exists(save_path) and not overwrite:
logger.error(f'File already exists at save_path: {save_path}')
raise Exception
json_dict = self.to_dict(strict=strict)
json.dump(json_dict, open(save_path, 'w'), indent=2, ensure_ascii=False)
def save_ckp(state,
checkpoint_dir,
best_model_dir,
filename_pth='checkpoint.pth'):
"""
checkpoint = {
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_ckp(checkpoint, is_best, checkpoint_dir, model_dir)
"""
make_dir_if_not_exists(checkpoint_dir)
make_dir_if_not_exists(best_model_dir)
best_fpath = os.path.join(best_model_dir, 'best_model.pth')
rename_f_path = os.path.join(best_model_dir, filename_pth)
f_path = os.path.join(checkpoint_dir, filename_pth)
if file_exists(best_fpath):
os.rename(best_fpath, rename_f_path)
move_file(rename_f_path, f_path)
torch.save(state, best_fpath)
def infer(self, input_type: Union[str, int],
output_type: Union[str, int],
input_path: Union[str, List[str]],
output_path: Union[str, List[str]],
show_max_score_only: bool = False,
show_class_label: bool = True,
show_class_label_score_only: bool = False,
show_keypoint_label: bool = True,
show_legends: bool = False,
show_bbox: bool = True,
show_keypoints: bool = True,
show_segmentation: bool = True,
color_bbox: list = None,
transparent_mask: bool = True,
transparency_alpha: float = 0.3,
ignore_keypoint_idx=None,
gt_path: Union[str, List[str]] = None,
result_json_path: str = None):
"""
Valid options for,
===
input_type:
---
["image", "image_list", "image_directory", "image_directories_list", "video",
"video_list", "video_directory" ]
output_type:
---
["return_summary", "show_image", "write_image", "write_video" ]
Returns
---
The inference result of all data formats.
"""
self.counter = infinite_sequence()
check_value(input_type,
check_from=["image", "image_list", "image_directory", "image_directories_list", "video",
"video_list", "video_directory"])
check_value(value=output_type, check_from=[
"show_image", "write_image", "write_video"])
self.gt_path = gt_path
if self.gt_path:
check_file_exists(gt_path)
with open(gt_path) as json_file:
self.gt_data = json.load(json_file)
if result_json_path is None:
if dir_exists(output_path):
self.result_json_path = f'{output_path}/result.json'
else:
_p = output_path.split('.')
_output_path = '.'.join(_p[:-1])
self.result_json_path = f'{_output_path}.json'
predict_image = partial(self._infer_image,
show_max_score_only=show_max_score_only,
show_class_label=show_class_label,
show_class_label_score_only=show_class_label_score_only,
show_keypoint_label=show_keypoint_label,
show_bbox=show_bbox, show_keypoints=show_keypoints, show_segmentation=show_segmentation,
color_bbox=color_bbox,
transparent_mask=transparent_mask, transparency_alpha=transparency_alpha,
ignore_keypoint_idx=ignore_keypoint_idx,
show_legends=show_legends
# gt_path=gt_path,
)
if input_type == "image":
if file_exists(input_path):
output = predict_image(input_path)
else:
raise Error
if output_type == "return_summary":
return self.pred_dataset
elif output_type == "show_image":
show_image(output)
elif output_type == "write_image":
cv2.imwrite(f'{output_path}', output)
elif output_type == "write_video":
raise NotImplementedError
else:
raise Exception
elif input_type == "image_list":
for image_path in tqdm(input_path, colour='#66cc66'):
output = predict_image(image_path)
if output_type == "show_image":
if show_image(output):
break
elif output_type == "write_image":
make_dir_if_not_exists(output_path)
filename = f.get_filename(image_path)
cv2.imwrite(f'{output_path}/{filename}', output)
elif output_type == "write_video":
raise NotImplementedError
else:
raise Exception
elif input_type == "image_directory":
image_path_list = f.dir_contents_path_list_with_extension(
dirpath=input_path,
extension=[".png", ".jpg", ".jpeg"])
for image_path in tqdm(image_path_list, colour='#66cc66'):
output = predict_image(image_path)
# output = self.draw_gt(gt_path, gt_data, image_path, output)
if output_type == "show_image":
if show_image(output):
break
elif output_type == "write_image":
make_dir_if_not_exists(output_path)
filename = f.get_filename(image_path)
cv2.imwrite(f'{output_path}/{filename}', output)
elif output_type == "write_video":
raise NotImplementedError
else:
raise Exception
elif input_type == "image_directories_list":
for image_directory in tqdm(input_path):
image_path_list = f.dir_contents_path_list_with_extension(
dirpath=input_path,
extension=[".png", ".jpg", ".jpeg"])
for image_path in tqdm(image_path_list, colour='#66cc66'):
output = predict_image(image_path)
if output_type == "show_image":
if show_image(output):
break
elif output_type == "write_image":
filename = f.get_filename(image_path)
directory_name = f.get_directory_name(image_path)
if f.dir_exists(f'{output_path}/{directory_name}'):
f.delete_all_files_in_dir(
f'{output_path}/{directory_name}')
else:
f.make_dir(f'{output_path}/{directory_name}')
cv2.imwrite(
f'{output_path}/{directory_name}/{filename}', output)
elif output_type == "write_video":
raise NotImplementedError
else:
raise Exception
elif input_type == "video":
raise NotImplementedError
elif input_type == "video_list":
raise NotImplementedError
elif input_type == "video_directory":
raise NotImplementedError
else:
raise Exception
# printj.cyan(self.pred_dataset)
self.write_predictions_json()
if self.gt_path:
# print(self.df)
# pip install openpyxl
self.df.to_excel(os.path.abspath(
f'{result_json_path}/../test_data.xlsx'))
def __init__(
self,
weights_path: str,
class_names: List[str] = None, num_classes: int = None,
keypoint_names: List[str] = None, num_keypoints: int = None,
model: str = "mask_rcnn_R_50_FPN_1x",
confidence_threshold: float = 0.5,
size_min: int = None,
size_max: int = None,
key_seg_together: bool = False,
gray_on=False,
crop_mode: int = None,
crop_mode2_rec: Union[int, List[int]] = None,
crop_mode3_sizes: Union[int, List[int]] = None,
crop_mode3_overlaps: Union[int, List[int]] = None,
detectron2_dir_path: str = "/home/jitesh/detectron/detectron2"
):
"""
D2Inferer
=========
Parameters:
------
- weights_path: str
- class_names: List[str] = None, num_classes: int = None,
- keypoint_names: List[str] = None, num_keypoints: int = None,
- model: str = "mask_rcnn_R_50_FPN_1x",
- confidence_threshold: float = 0.5,
- size_min: int = None,
- size_max: int = None,
- key_seg_together: bool = False,
- detectron2_dir_path: str = "/home/jitesh/detectron/detectron2"
- crop_mode = 1 : crop between points (0, 0) and (a, a), where a is min(height, width)
- crop_mode = 2 : crop between points crop_rec[0] and crop_rec[1], crop_rec is defined by the user through parameter
- crop_mode = 3 : chop, infer and merge
- crop_rec : list of points of rectangle to crop
"""
self.df = pd.DataFrame(data=[], columns=[])
self.gray_on = gray_on
self.crop_mode = crop_mode
self.crop_rec = crop_mode2_rec
self.crop_mode3_sizes = crop_mode3_sizes
self.crop_mode3_overlaps = crop_mode3_overlaps
if class_names is None:
class_names = ['']
if keypoint_names is None:
keypoint_names = ['']
self.key_seg_together = key_seg_together
self.weights_path = weights_path
self.class_names = class_names
if num_classes is None:
self.num_classes = len(class_names)
else:
assert num_classes == len(class_names)
self.num_classes = num_classes
self.keypoint_names = keypoint_names
if num_keypoints is None:
self.num_keypoints = len(keypoint_names)
else:
assert num_keypoints == len(keypoint_names)
self.num_keypoints = num_keypoints
self.confidence_threshold = confidence_threshold
self.model = model
if "COCO-Detection" in self.model:
self.model = self.model
elif "COCO-Keypoints" in self.model:
self.model = self.model
elif "COCO-InstanceSegmentation" in self.model:
self.model = self.model
elif "COCO-PanopticSegmentation" in self.model:
self.model = self.model
elif "LVIS-InstanceSegmentation" in self.model:
self.model = self.model
elif "Misc" in model:
self.model = model
# train_type = 'seg'
elif "rpn" in model:
self.model = "COCO-Detection/" + model
elif "keypoint" in model:
self.model = "COCO-Keypoints/" + model
elif "mask" in model:
self.model = "COCO-InstanceSegmentation/" + model
else:
printj.red.bold_on_black(f'{model} is not in the dictionary.\
Choose the correct model.')
raise Exception
if ".yaml" in self.model:
self.model = self.model
else:
self.model = self.model + ".yaml"
self.cfg = get_cfg()
model_conf_path = f"{detectron2_dir_path}/configs/{self.model}"
if not file_exists(model_conf_path):
printj.red(f"Invalid model: {model}\nOr")
printj.red(f"File not found: {model_conf_path}")
raise Exception
self.cfg.merge_from_file(model_conf_path)
self.cfg.MODEL.WEIGHTS = self.weights_path
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = self.num_classes
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.confidence_threshold
self.cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = self.num_keypoints
if "mask" in self.model.lower() or "segmentation" in self.model.lower():
self.cfg.MODEL.MASK_ON = True
# self.cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT=0.5
if size_min is not None:
self.cfg.INPUT.MIN_SIZE_TRAIN = size_min
self.cfg.INPUT.MIN_SIZE_TEST = size_min
if size_max is not None:
self.cfg.INPUT.MAX_SIZE_TRAIN = size_max
self.cfg.INPUT.MAX_SIZE_TEST = size_max
self.predictor = DefaultPredictor(self.cfg)
self.pred_dataset = []
self.palette = np.array(color_palette(
palette='hls', n_colors=self.num_classes+1))*255
def move_images(
self, dst_img_dir: str,
preserve_filenames: bool=False, overwrite_duplicates: bool=False, update_img_paths: bool=True, overwrite: bool=False,
show_pbar: bool=True
):
"""
Combines all image directories specified in the coco_url of each coco image in self.images
to a single image directory.
dst_img_dir: The directory where you would like to save the combined image set.
preserve_filenames: If False, unique filenames will be generated so as to not create a filename conflict.
overwrite_duplicates: Only applicable when preserve_filenames=True.
In the event that two images with the same filename are moved to dst_img_dir from
two different source folders, an error will be raised if overwrite_duplicates=False.
If overwrite_duplicates=True, the second copy will be overwrite the first copy.
update_img_paths: If True, all coco_url paths specified in self.images will be updated to reflect the new
combined image directory.
overwrite: If True, all files in dst_img_dir will be deleted before copying images into the folder.
"""
used_img_dir_list = []
for coco_image in self.images:
used_img_dir = get_dirpath_from_filepath(coco_image.coco_url)
if used_img_dir not in used_img_dir_list:
check_dir_exists(used_img_dir)
used_img_dir_list.append(used_img_dir)
if len(used_img_dir_list) == 0:
logger.error(f"Couldn't parse used_img_dir_list.")
logger.error(f"Are the coco_url paths in your dataset's image dictionary correct?")
raise Exception
make_dir_if_not_exists(dst_img_dir)
if get_dir_contents_len(dst_img_dir) > 0:
if overwrite:
delete_all_files_in_dir(dst_img_dir, ask_permission=False)
else:
logger.error(f'dst_img_dir={dst_img_dir} is not empty.')
logger.error('Please use overwrite=True if you would like to delete the contents before proceeding.')
raise Exception
pbar = tqdm(total=len(self.images), unit='image(s)') if show_pbar else None
pbar.set_description(f'Moving Images...')
for coco_image in self.images:
if not preserve_filenames:
img_extension = get_extension_from_path(coco_image.coco_url)
dst_img_path = get_next_dump_path(
dump_dir=dst_img_dir,
file_extension=img_extension
)
dst_img_path = rel_to_abs_path(dst_img_path)
else:
img_filename = get_filename(coco_image.coco_url)
dst_img_path = f'{dst_img_dir}/{img_filename}'
if file_exists(dst_img_path) and not overwrite_duplicates:
logger.error(f'Failed to copy {coco_image.coco_url} to {dst_img_dir}')
logger.error(f'{img_filename} already exists in destination directory.')
logger.error(f'Hint: In order to use preserve_filenames=True, all filenames in the dataset must be unique.')
logger.error(
f'Suggestion: Either update the filenames to be unique or use preserve_filenames=False' + \
f' in order to automatically assign the destination filename.'
)
raise Exception
copy_file(src_path=coco_image.coco_url, dest_path=dst_img_path, silent=True)
if update_img_paths:
coco_image.coco_url = dst_img_path
coco_image.file_name = get_filename(dst_img_path)
if pbar is not None:
pbar.update(1)
if pbar is not None:
pbar.close()
def __init__(
self,
coco_ann_path: str,
img_path: str,
val_coco_ann_path: str,
val_img_path: str,
output_dir_path: str,
resume: bool = True,
class_names: List[str] = None,
num_classes: int = None,
keypoint_names: List[str] = None,
num_keypoints: int = None,
model: str = "mask_rcnn_R_50_FPN_1x",
instance_train: str = "training_instance1",
min_size_train: int = None,
max_size_train: int = None,
min_size_test: int = None,
max_size_test: int = None,
max_iter: int = 10000,
batch_size_per_image: int = 512,
checkpoint_period: int = None,
score_thresh: int = None,
key_seg_together: bool = False,
aug_on: bool = True,
train_val: bool = False,
aug_settings_file_path: str = None,
aug_vis_save_path: str = 'aug_vis.png',
show_aug_seg: bool = False,
aug_n_rows: int = 3,
aug_n_cols: int = 5,
aug_save_dims: Tuple[int] = (3 * 500, 5 * 500),
device: str = 'cuda',
num_workers: int = 2,
images_per_batch: int = 2,
base_lr: float = 0.003,
decrease_lr_by_ratio: float = 0.1,
lr_steps: tuple = (30000, ),
detectron2_dir_path: str = None,
val_on: bool = False,
instance_test: str = "test_instance1",
val_eval_period: int = 100,
vis_period: int = 0,
train_type: str = None,
):
"""
D2Trainer
=========
Parameters:
------
output_dir_path: str
class_names: List[str] = None, num_classes: int = None,
keypoint_names: List[str] = None, num_keypoints: int = None,
model: str = "mask_rcnn_R_50_FPN_1x",
confidence_threshold: float = 0.5,
min_size_train: int = None,
max_size_train: int = None,
key_seg_together: bool = False,
detectron2_dir_path: str = "/home/jitesh/detectron/detectron2"
"""
self.key_seg_together = key_seg_together
self.coco_ann_path = coco_ann_path
self.img_path = img_path
self.val_coco_ann_path = val_coco_ann_path
self.val_img_path = val_img_path
self.output_dir_path = output_dir_path
self.instance_train = instance_train
self.resume = resume
self.device = device
self.num_workers = num_workers
self.images_per_batch = images_per_batch
self.batch_size_per_image = batch_size_per_image
self.checkpoint_period = checkpoint_period
self.score_thresh = score_thresh
self.base_lr = base_lr
self.decrease_lr_by_ratio = decrease_lr_by_ratio
self.lr_steps = lr_steps
self.max_iter = max_iter
self.val_on = val_on
self.instance_test = instance_test
self.val_eval_period = val_eval_period
self.vis_period = vis_period
""" Load annotations json """
with open(self.coco_ann_path) as json_file:
self.coco_ann_data = json.load(json_file)
self.categories = self.coco_ann_data["categories"]
if class_names is None:
# self.class_names = ['']
self.class_names = [
category["name"] for category in self.categories
]
else:
self.class_names = class_names
if num_classes is None:
self.num_classes = len(self.class_names)
else:
printj.red(f'num_classes: {num_classes}')
printj.red(f'len(self.class_names): {len(self.class_names)}')
assert num_classes == len(self.class_names)
self.num_classes = num_classes
if keypoint_names is None:
self.keypoint_names = ['']
else:
self.keypoint_names = keypoint_names
if num_keypoints is None:
if keypoint_names == ['']:
self.num_keypoints = 0
else:
self.num_keypoints = len(self.keypoint_names)
else:
assert num_keypoints == len(self.keypoint_names)
self.num_keypoints = num_keypoints
self.model = model
if "COCO-Detection" in self.model:
self.model = self.model
train_type = 'bbox'
elif "COCO-Keypoints" in self.model:
self.model = self.model
train_type = 'kpt'
elif "COCO-InstanceSegmentation" in self.model:
self.model = self.model
train_type = 'seg'
elif "COCO-PanopticSegmentation" in self.model:
self.model = self.model
train_type = 'seg'
elif "LVIS-InstanceSegmentation" in self.model:
self.model = self.model
train_type = 'seg'
elif "Misc" in model:
self.model = model
train_type = 'seg'
elif "rpn" in model or "fast" in model:
self.model = "COCO-Detection/" + model
train_type = 'bbox'
elif "keypoint" in model:
self.model = "COCO-Keypoints/" + model
train_type = 'kpt'
elif "mask" in model:
self.model = "COCO-InstanceSegmentation/" + model
train_type = 'seg'
elif train_type:
self.model = model
train_type = train_type
else:
printj.red.bold_on_black(f'{model} is not in the dictionary.\
Choose the correct model.')
raise Exception
if ".yaml" in self.model:
self.model = self.model
else:
self.model = self.model + ".yaml"
if detectron2_dir_path:
model_conf_path = f"{detectron2_dir_path}/configs/{self.model}"
else:
model_conf_path = model_zoo.get_config_file(self.model)
if not file_exists(model_conf_path):
printj.red(f"Invalid model: {model}\nOr")
printj.red(f"File not found: {model_conf_path}")
raise Exception
""" register """
register_coco_instances(name=self.instance_train,
metadata={},
json_file=self.coco_ann_path,
image_root=self.img_path)
MetadataCatalog.get(
self.instance_train).thing_classes = self.class_names
# sys.exit(self.class_names)
if val_on:
register_coco_instances(name=self.instance_test,
metadata={},
json_file=self.val_coco_ann_path,
image_root=self.val_img_path)
MetadataCatalog.get(
self.instance_test).thing_classes = self.class_names
""" cfg """
self.cfg = get_cfg()
self.cfg.merge_from_file(model_conf_path)
self.cfg.DATASETS.TRAIN = tuple([self.instance_train])
self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(self.model)
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = self.num_classes
self.cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = self.num_keypoints
self.cfg.DATALOADER.NUM_WORKERS = self.num_workers
self.cfg.SOLVER.IMS_PER_BATCH = self.images_per_batch
self.cfg.SOLVER.BASE_LR = self.base_lr
self.cfg.MODEL.DEVICE = self.device
self.cfg.OUTPUT_DIR = self.output_dir_path
if self.lr_steps:
self.cfg.SOLVER.GAMMA = self.decrease_lr_by_ratio
self.cfg.SOLVER.STEPS = self.lr_steps
if self.max_iter:
self.cfg.SOLVER.MAX_ITER = self.max_iter
if self.batch_size_per_image:
self.cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = self.batch_size_per_image
if self.checkpoint_period:
self.cfg.SOLVER.CHECKPOINT_PERIOD = self.checkpoint_period
if self.vis_period:
self.cfg.VIS_PERIOD = self.vis_period
if score_thresh:
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thresh
if self.val_on:
self.cfg.DATASETS.TEST = tuple([self.instance_test])
self.cfg.TEST.EVAL_PERIOD = self.val_eval_period
make_dir_if_not_exists(self.cfg.OUTPUT_DIR)
if not self.resume:
delete_dir_if_exists(self.cfg.OUTPUT_DIR)
make_dir_if_not_exists(self.cfg.OUTPUT_DIR)
if "mask" in self.model.lower() or "segmentation" in self.model.lower(
):
self.cfg.MODEL.MASK_ON = True
else:
self.cfg.MODEL.MASK_ON = False
# self.cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT=0.5
# Train Size Parameters
if min_size_train is not None:
self.cfg.INPUT.MIN_SIZE_TRAIN = min_size_train
if max_size_train is not None:
self.cfg.INPUT.MAX_SIZE_TRAIN = max_size_train
# Test Size Parameters
if min_size_test is not None:
self.cfg.INPUT.MIN_SIZE_TEST = min_size_test
elif min_size_train is not None:
self.cfg.INPUT.MIN_SIZE_TEST = min_size_train
if max_size_test is not None:
self.cfg.INPUT.MAX_SIZE_TEST = max_size_test
elif max_size_train is not None:
self.cfg.INPUT.MAX_SIZE_TEST = max_size_train
self.cfg.INPUT.MIN_SIZE_TEST = min_size_train
""" def train() """
self.aug_settings_file_path = aug_settings_file_path
self.aug_on = aug_on
self.train_val = train_val
self.train_type = train_type
self.aug_vis_save_path = aug_vis_save_path
self.show_aug_seg = show_aug_seg
self.aug_n_rows = aug_n_rows
self.aug_n_cols = aug_n_cols
self.aug_save_dims = aug_save_dims