def _draw(self, img: np.ndarray, bbox: BBox, kpt_list: list,
kpt_skeleton: list, kpt_label_list: list, seg: Segmentation,
cat_name: str) -> np.ndarray:
for render_target in self.render_order:
if render_target.lower() in ['seg', 'segmentation']:
if self.show_seg:
img = draw_segmentation(img=img,
segmentation=seg,
color=self.seg_color,
transparent=self.transparent_seg)
elif render_target.lower() in ['skel', 'skeleton']:
if self.show_skeleton:
img = draw_skeleton(
img=img,
keypoints=kpt_list,
keypoint_skeleton=kpt_skeleton,
index_offset=self.skeleton_index_offset,
thickness=self.skeleton_thickness,
color=self.skeleton_color)
elif render_target.lower() in [
'kpt', 'kpts', 'keypoint', 'keypoints'
]:
if self.show_kpts:
if self.kpt_label_thickness is not None:
kpt_label_thickness = self.kpt_label_thickness
else:
kpt_label_thickness = self.kpt_radius - 1 if self.kpt_radius > 1 else 1
img = draw_keypoints(
img=img,
keypoints=kpt_list,
radius=self.kpt_radius,
color=self.kpt_color,
keypoint_labels=kpt_label_list,
show_keypoints_labels=self.show_kpt_labels,
label_thickness=kpt_label_thickness,
label_only=self.show_kpt_labels_only)
elif render_target.lower() in [
'bbox', 'bounding_box', 'bounding box'
]:
if self.show_bbox:
bbox_text = cat_name if self.show_bbox_labels or self.show_bbox_labels_only else None
bbox_label_thickness = self.bbox_label_thickness if self.bbox_label_thickness is not None else self.bbox_thickness
img = draw_bbox(img=img,
bbox=bbox,
color=self.bbox_color,
thickness=self.bbox_thickness,
text=bbox_text,
label_thickness=bbox_label_thickness,
label_only=self.show_bbox_labels_only)
else:
logger.error(f"Invalid render_target: {render_target}")
logger.error(f"Options: seg, skeleton, kpt, bbox")
raise Exception
return img
def mapper(dataset_dict):
# Implement a mapper, similar to the default DatasetMapper, but with your own customizations
dataset_dict = copy.deepcopy(dataset_dict) # it will be modified by code below
image = utils.read_image(dataset_dict["file_name"], format="BGR")
#print(dataset_dict["file_name"])
#cv2.imwrite(os.path.join(aug_save_path, f'{str(count.count)}_asd.jpg'),image)
handler = load_augmentation_settings(handler_save_path = 'test_handler.json')
for i in range(len(dataset_dict["annotations"])):
dataset_dict["annotations"][i]["segmentation"] = []
image, dataset_dict = handler(image=image, dataset_dict_detectron=dataset_dict)
annots = []
for item in dataset_dict["annotations"]:
annots.append(item)
dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
instances = utils.annotations_to_instances(annots, image.shape[:2])
dataset_dict["instances"] = utils.filter_empty_instances(instances)
if True:
vis_img = image.copy()
bbox_list = [BBox.from_list(vals) for vals in dataset_dict["instances"].gt_boxes.tensor.numpy().tolist()]
# seg_list = [Segmentation([Polygon.from_list(poly.tolist(), demarcation=False) for poly in seg_polys]) for seg_polys in dataset_dict["instances"].gt_masks.polygons]
for bbox in (bbox_list):
# if len(seg) > 0 and False:
# vis_img = draw_segmentation(img=vis_img, segmentation=seg, transparent=True)
vis_img = draw_bbox(img=vis_img, bbox=bbox)
height,width,channels = vis_img.shape
print(height,width)
if count.count<100:
cv2.imwrite(os.path.join(aug_save_path, f'{str(count.count)}.jpg'),vis_img)
count.count_add()
aug_vis.step(vis_img)
return dataset_dict
def __call__(self, dataset_dict):
"""
Args:
dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
Returns:
dict: a format that builtin models in detectron2 accept
"""
dataset_dict = copy.deepcopy(dataset_dict) # it will be modified by code below
# USER: Write your own image loading if it's not from a file
image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
utils.check_image_size(dataset_dict, image)
for i in range(len(dataset_dict["annotations"])):
dataset_dict["annotations"][i]["segmentation"] = []
### my code ##
image, dataset_dict = self.aug_handler(image=image, dataset_dict_detectron=dataset_dict)
### my code ##
if "annotations" not in dataset_dict:
image, transforms = T.apply_transform_gens(
([self.crop_gen] if self.crop_gen else []) + self.tfm_gens, image
)
else:
# Crop around an instance if there are instances in the image.
# USER: Remove if you don't use cropping
if self.crop_gen:
crop_tfm = utils.gen_crop_transform_with_instance(
self.crop_gen.get_crop_size(image.shape[:2]),
image.shape[:2],
np.random.choice(dataset_dict["annotations"]),
)
image = crop_tfm.apply_image(image)
image, transforms = T.apply_transform_gens(self.tfm_gens, image)
if self.crop_gen:
transforms = crop_tfm + transforms
image_shape = image.shape[:2] # h, w
# Pytorch's dataloader is efficient on torch.Tensor due to shared-memory,
# but not efficient on large generic data structures due to the use of pickle & mp.Queue.
# Therefore it's important to use torch.Tensor.
dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
# USER: Remove if you don't use pre-computed proposals.
if self.load_proposals:
utils.transform_proposals(
dataset_dict, image_shape, transforms, self.min_box_side_len, self.proposal_topk
)
if not self.is_train:
# USER: Modify this if you want to keep them for some reason.
dataset_dict.pop("annotations", None)
dataset_dict.pop("sem_seg_file_name", None)
return dataset_dict
if "annotations" in dataset_dict:
# USER: Modify this if you want to keep them for some reason.
for anno in dataset_dict["annotations"]:
if not self.mask_on:
anno.pop("segmentation", None)
if not self.keypoint_on:
anno.pop("keypoints", None)
# USER: Implement additional transformations if you have other types of data
annos = [
utils.transform_instance_annotations(
obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices
)
for obj in dataset_dict.pop("annotations")
if obj.get("iscrowd", 0) == 0
]
instances = utils.annotations_to_instances(
annos, image_shape, mask_format=self.mask_format
)
# Create a tight bounding box from masks, useful when image is cropped
if self.crop_gen and instances.has("gt_masks"):
instances.gt_boxes = instances.gt_masks.get_bounding_boxes()
dataset_dict["instances"] = utils.filter_empty_instances(instances)
# USER: Remove if you don't do semantic/panoptic segmentation.
if "sem_seg_file_name" in dataset_dict:
with PathManager.open(dataset_dict.pop("sem_seg_file_name"), "rb") as f:
sem_seg_gt = Image.open(f)
sem_seg_gt = np.asarray(sem_seg_gt, dtype="uint8")
sem_seg_gt = transforms.apply_segmentation(sem_seg_gt)
sem_seg_gt = torch.as_tensor(sem_seg_gt.astype("long"))
dataset_dict["sem_seg"] = sem_seg_gt
#### this uses local variable, use with caution ######
if True:
vis_img = image.copy()
bbox_list = [BBox.from_list(vals) for vals in dataset_dict["instances"].gt_boxes.tensor.numpy().tolist()]
# seg_list = [Segmentation([Polygon.from_list(poly.tolist(), demarcation=False) for poly in seg_polys]) for seg_polys in dataset_dict["instances"].gt_masks.polygons]
for bbox in (bbox_list):
# if len(seg) > 0 and False:
# vis_img = draw_segmentation(img=vis_img, segmentation=seg, transparent=True)
vis_img = draw_bbox(img=vis_img, bbox=bbox)
aug_vis.step(vis_img)
return dataset_dict
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 infer(
path: str,
weights_path: str,
thresh: int = 0.5,
key: str = 'R',
infer_dump_dir: str = '',
model: str = 'mask_rcnn_R_50_FPN_1x',
size: int = 1024,
class_names: List[str] = ['hook'],
gt_path:
str = '/home/jitesh/3d/data/coco_data/hook_test/json/cropped_hook.json'):
# class_names=['hook', 'pole']
# class_names=['hook']
conf_thresh = 0.001
show_bbox_border = True
gt_dataset = COCO_Dataset.load_from_path(json_path=gt_path)
inferer_seg = inferer(
weights_path=weights_path,
confidence_threshold=0.1,
# num_classes=1,
# num_classes=2,
class_names=class_names,
# class_names=['hook'],
model='keypoint_rcnn_R_50_FPN_1x',
# model='faster_rcnn_X_101_32x8d_FPN_3x',
# model='faster_rcnn_R_101_FPN_3x',
# model=model,
)
inferer_seg.cfg.INPUT.MIN_SIZE_TEST = size
inferer_seg.cfg.INPUT.MAX_SIZE_TEST = size
inferer_seg.cfg.MODEL.MASK_ON = True
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data7/weights/Keypoints_R_50_1x_aug_cm_seg_val_1/model_0009999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data7_0.1/weights/Keypoints_R_50_1x_aug_cm_seg_val_3/model_0009999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data7_0.1/weights/Keypoints_R_50_1x_aug_cm_seg_val_1/model_0007999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data8/weights/Keypoints_R_50_1x_aug_key_seg_val_1/model_0009999.pth'
weights_path = '/home/jitesh/3d/data/coco_data/hook_sim_real_data8/weights/Keypoints_R_50_1x_aug_key_seg_val_2/model_0004999.pth'
# inferer_key = jDetectron2KeypointInferer(
# weights_path=weights_path,
# # ref_coco_ann_path=f'/home/jitesh/3d/data/coco_data/hook_real1/json/hook.json',
# # categories_path=f'/home/jitesh/3d/data/categories/hook_infer.json',
# # categories_path=f'/home/jitesh/3d/data/categories/hook_7ckpt.json',
# categories_path=f'/home/jitesh/3d/data/categories/hook_7ckpt_pole.json',
# target_category='hook',
# model_name='keypoint_rcnn_R_50_FPN_1x',
# bbox_threshold=bbox_thresh,
# kpt_threshold=kpt_thresh,
# key_box='hook',
# )
# k_size = 1024
# inferer_key.cfg.INPUT.MIN_SIZE_TEST = k_size
# inferer_key.cfg.INPUT.MAX_SIZE_TEST = k_size
possible_modes = ['save', 'preview']
mode = 'save'
check_value(mode, valid_value_list=possible_modes)
# make_dir_if_not_exists(infer_dump_dir)
img_extensions = ['jpg', 'JPG', 'png', 'PNG']
img_pathlist = get_all_files_in_extension_list(
dir_path=f'{path}', extension_list=img_extensions)
img_pathlist.sort()
confirm_folder(infer_dump_dir, mode)
# confirm_folder(f'{infer_dump_dir}/good_seg', mode)
# confirm_folder(f'{infer_dump_dir}/good_cropped', mode)
# confirm_folder(f'{infer_dump_dir}/good', mode)
# confirm_folder(f'{infer_dump_dir}/G(>4D) P(>4D)', mode)
# confirm_folder(f'{infer_dump_dir}/G(>4D) P(<4D)', mode)
# confirm_folder(f'{infer_dump_dir}/G(<4D) P(>4D)', mode)
# confirm_folder(f'{infer_dump_dir}/G(<4D) P(<4D)', mode)
# confirm_folder(f'{infer_dump_dir}/bad', mode)
confirm_folder(f'{infer_dump_dir}/infer_key_seg', mode)
count = 0
start = datetime.now()
df = pd.DataFrame(data=[],
columns=[
'gt_d',
'pred_d',
'gt_ab',
'pred_ab',
'gt_ratio',
'pred_ratio',
'gt_ratio>4',
'pred_ratio>4',
'correct_above4d_ratio',
'incorrect_above4d_ratio',
'correct_below4d_ratio',
'incorrect_below4d_ratio',
])
# 'image_path'])
for i, img_path in enumerate(tqdm(
img_pathlist,
desc='Writing images',
)):
img_filename = get_filename(img_path)
# if not '201005_70_縮小革命PB020261.jpg' in img_path:
# continue
# if i > 19:
# continue
printj.purple(img_path)
img = cv2.imread(img_path)
result = img
# print(f'shape {img.shape}')
# cv2.imshow('i', img)
# cv2.waitKey(100000)
# continue
score_list, pred_class_list, bbox_list, pred_masks_list, pred_keypoints_list, vis_keypoints_list, kpt_confidences_list = inferer_seg.predict(
img=img)
# printj.blue(pred_masks_list)
max_hook_score = -1
max_pole_score = -1
diameter = -1
len_ab = -1
found_hook = False
found_pole = False
for score, pred_class, bbox, mask, keypoints, vis_keypoints, kpt_confidences in zip(
score_list, pred_class_list, bbox_list, pred_masks_list,
pred_keypoints_list, vis_keypoints_list, kpt_confidences_list):
if pred_class == 'pole':
found_pole = True
if max_pole_score < score:
# if True:
max_pole_score = score
diameter = compute_diameter(mask)
# result = draw_bool_mask(img=result, mask=mask, color=[
# 0, 255, 255],
# transparent=True
# )
pole_bbox_text = f'pole {str(round(score, 2))}'
pole_bbox = bbox
pole_mask = mask
# result = draw_bbox(img=result, bbox=bbox,
# text=pole_bbox_text, label_only=not show_bbox_border, label_orientation='bottom')
printj.blue(f'diameter={diameter}')
if pred_class == 'hook':
# printj.green.bold_on_yellow(score)
found_hook = True
if max_hook_score < score:
# if True:
max_hook_score = score
hook_bbox = BBox.buffer(bbox)
hook_score = round(score, 2)
hook_mask = mask
hook_keypoints = keypoints
hook_vis_keypoints = vis_keypoints
hook_kpt_confidences = kpt_confidences
# xmin, ymin, xmax, ymax = bbox.to_int().to_list()
# _xmin, _ymin, _xmax, _ymax = _bbox.to_int().to_list()
# width = _xmax-_xmin
# height = _ymax-_ymin
# scale = 0.2
# xmin = max(int(_xmin - width*scale), 0)
# xmax = min(int(_xmax + width*scale), img.shape[1])
# ymin = max(int(_ymin - height*scale), 0)
# ymax = min(int(_ymax + height*scale), img.shape[0])
# printj.red(score)
# printj.red(bbox)
# return
# img = draw_bbox(img=img, bbox=_bbox, color=[
# 0, 255, 255], thickness=2, text=f"{pred_class} {round(score, 3)}",
# label_orientation='top')
# img = draw_bbox(img=img, bbox=_bbox, color=[
# 0, 255, 255], thickness=2, text=f"{pred_class} {round(score, 3)}",
# label_orientation='bottom')
# result = draw_bool_mask(img=result, mask=mask, color=[
# 255, 255, 0],
# transparent=True
# )
# result = result
# bbox_text = str(round(score, 4))
# result = draw_bbox(img=result, bbox=bbox,
# text=bbox_text, label_only=not show_bbox_border)
bbox_label_mode = 'euler'
# result = draw_keypoints(
# img=result, keypoints=vis_keypoints, radius=2, color=[0, 0, 255],
# # keypoint_labels=kpt_labels, show_keypoints_labels=True, label_thickness=1,
# # ignore_kpt_idx=conf_idx_list
# )
kpt_labels = [
"kpt-a", "kpt-b", "kpt-cb", "kpt-c", "kpt-cd", "kpt-d",
"kpt-e"
]
kpt_skeleton = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5],
[5, 6]]
conf_idx_list = np.argwhere(
np.array(kpt_confidences) > conf_thresh).reshape(-1)
not_conf_idx_list = np.argwhere(
np.array(kpt_confidences) <= conf_thresh).reshape(
-1).astype(int)
conf_keypoints, conf_kpt_labels = np.array(vis_keypoints)[
conf_idx_list], np.array(kpt_labels)[conf_idx_list]
not_conf_keypoints, not_conf_kpt_labels = np.array(
vis_keypoints)[not_conf_idx_list], np.array(
kpt_labels)[not_conf_idx_list]
cleaned_keypoints = np.array(vis_keypoints.copy()).astype(
np.float32)
# result = draw_bool_mask(img=result, mask=mask, color=[
# 255, 255, 0],
# transparent=True
# )
# result, len_ab = draw_inference_on_hook2(img=result, cleaned_keypoints=cleaned_keypoints, kpt_labels=kpt_labels, kpt_skeleton=kpt_skeleton,
# score=score, bbox=_bbox, vis_keypoints=vis_keypoints, kpt_confidences=kpt_confidences, conf_idx_list=conf_idx_list, not_conf_idx_list=not_conf_idx_list,
# conf_keypoints=conf_keypoints, conf_kpt_labels=conf_kpt_labels, not_conf_keypoints=not_conf_keypoints, not_conf_kpt_labels=not_conf_kpt_labels,
# conf_thresh=conf_thresh, show_bbox_border=show_bbox_border, bbox_label_mode=bbox_label_mode, index_offset=0, diameter=diameter)
# result=result
# printj.green(_bbox)
# printj.green(_bbox.to_int())
# printj.green(_bbox.to_int().to_list())
printj.green.on_white(max_hook_score)
if found_pole:
result = draw_bool_mask(img=result,
mask=pole_mask,
color=[0, 255, 255],
transparent=True)
result = draw_bbox(img=result,
bbox=pole_bbox,
text=pole_bbox_text,
label_only=not show_bbox_border,
label_orientation='top')
result = draw_bbox(img=result,
bbox=pole_bbox,
text=pole_bbox_text,
label_only=not show_bbox_border,
label_orientation='bottom')
if found_hook:
result = draw_bool_mask(img=result,
mask=hook_mask,
color=[255, 255, 0],
transparent=True)
result, len_ab = draw_inference_on_hook2(
img=result,
cleaned_keypoints=cleaned_keypoints,
kpt_labels=kpt_labels,
kpt_skeleton=kpt_skeleton,
score=hook_score,
bbox=hook_bbox,
vis_keypoints=hook_vis_keypoints,
kpt_confidences=hook_kpt_confidences,
conf_idx_list=conf_idx_list,
not_conf_idx_list=not_conf_idx_list,
conf_keypoints=conf_keypoints,
conf_kpt_labels=conf_kpt_labels,
not_conf_keypoints=not_conf_keypoints,
not_conf_kpt_labels=not_conf_kpt_labels,
conf_thresh=conf_thresh,
show_bbox_border=show_bbox_border,
bbox_label_mode=bbox_label_mode,
index_offset=0,
diameter=diameter)
printj.purple(len_ab)
if len_ab == 0:
printj.green(keypoints)
result = draw_info_box(result, len_ab, diameter)
# img: np.ndarray, cleaned_keypoints, kpt_labels: List[str], kpt_skeleton: List[list],
# score: float, bbox: BBox, vis_keypoints: list, kpt_confidences: list, conf_idx_list: list, not_conf_idx_list: list,
# conf_keypoints, conf_kpt_labels, not_conf_keypoints, not_conf_kpt_labels,
# conf_thresh: float = 0.3, show_bbox_border: bool = False, bbox_label_mode: str = 'euler', index_offset: int = 0, diameter=1
# cv2.imshow('i', result)
# # cv2.imwrite('i', result)
# cv2.waitKey(10000)
# quit_flag = cv_simple_image_viewer(img=result, preview_width=1000)
# if quit_flag:
# break
# cv2.imwrite(f"{infer_dump_dir}/good_seg/{img_filename}", result)
cv2.imwrite(f"{infer_dump_dir}/infer_key_seg/{img_filename}", result)