""" This file computes the IoU-based and centroid-distance-based metrics in a symmetric manner\n """ import sys, os sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) import numpy as np from copy import deepcopy from objects.dict_utils import ( get_base_part_idx, get_bbox_vertices, remove_handles, compute_overall_bbox_size, rescale_object, find_part_mapping, zero_center_object, ) from objects.motions import transform_all_parts from metrics.giou import sampling_giou, sampling_cDist def _get_scores( src_dict, tgt_dict, original_src_bbox_vertices, original_tgt_bbox_vertices, mapping, num_states, rotation_fix_range, num_samples, iou_include_base, ): # Record the indices of the base parts of the src objects src_base_idx = get_base_part_idx(src_dict) # Compute the sum of IoU between the generated object and the candidate object over a number of articulation states num_parts_in_src = len(src_dict["diffuse_tree"]) iou_per_part_and_state = np.zeros((num_parts_in_src, num_states), dtype=np.float32) cDist_per_part_and_state = np.zeros( (num_parts_in_src, num_states), dtype=np.float32 ) states = np.linspace(0, 1, num_states) for state_idx, state in enumerate(states): # Get a fresh copy of the bounding box vertices in rest pose src_bbox_vertices = deepcopy(original_src_bbox_vertices) tgt_bbox_vertices = deepcopy(original_tgt_bbox_vertices) # Transform the objects to the current state using the joints src_part_transfomrations = transform_all_parts( src_bbox_vertices, src_dict, state, rotation_fix_range=rotation_fix_range, ) tgt_part_transfomrations = transform_all_parts( tgt_bbox_vertices, tgt_dict, state, rotation_fix_range=rotation_fix_range, ) # Compute the IoU between the two objects using the transformed bounding boxes and the part mapping for src_part_idx in range(num_parts_in_src): # Get the index of the corresponding part in the candidate object tgt_part_idx = int(mapping[src_part_idx, 0]) # Always use a fresh copy of the bounding box vertices in rest pose in case dry_run=False is incorrectly set src_part_bbox_vertices = deepcopy(original_src_bbox_vertices)[src_part_idx] tgt_part_bbox_vertices = deepcopy(original_tgt_bbox_vertices)[tgt_part_idx] # Compute the sampling-based IoU between the two parts iou_per_part_and_state[src_part_idx, state_idx] = sampling_giou( src_part_bbox_vertices, tgt_part_bbox_vertices, src_part_transfomrations[src_part_idx], tgt_part_transfomrations[tgt_part_idx], num_samples=num_samples, ) # Compute the centriod distance between the two matched parts cDist_per_part_and_state[src_part_idx, state_idx] = sampling_cDist( src_dict["diffuse_tree"][src_part_idx], tgt_dict["diffuse_tree"][tgt_part_idx], src_part_transfomrations[src_part_idx], tgt_part_transfomrations[tgt_part_idx], ) # IoU and cDist at the resting state per_part_iou_avg_at_rest = iou_per_part_and_state[:, 0] per_part_cDist_avg_at_rest = cDist_per_part_and_state[:, 0] # Average the IoU over the states per_part_iou_avg_over_states = np.sum(iou_per_part_and_state, axis=1) / num_states # Average the cDist over the states per_part_cDist_avg_over_states = ( np.sum(cDist_per_part_and_state, axis=1) / num_states ) # Remove the base part if specified if not iou_include_base: per_part_iou_avg_over_states = np.delete( per_part_iou_avg_over_states, src_base_idx ) per_part_iou_avg_at_rest = np.delete(per_part_iou_avg_at_rest, src_base_idx) per_part_cDist_avg_over_states = np.delete( per_part_cDist_avg_over_states, src_base_idx ) per_part_cDist_avg_at_rest = np.delete(per_part_cDist_avg_at_rest, src_base_idx) aid_iou = float(np.mean(per_part_iou_avg_over_states)) if len(per_part_iou_avg_over_states) > 0 else 0 aid_cdist = float(np.mean(per_part_cDist_avg_over_states)) if len(per_part_cDist_avg_over_states) > 0 else 1 rid_iou = float(np.mean(per_part_iou_avg_at_rest)) if len(per_part_iou_avg_at_rest) > 0 else 0 rid_cdist = float(np.mean(per_part_cDist_avg_at_rest)) if len(per_part_cDist_avg_at_rest) > 0 else 1 return { "AS-IoU": 1. - aid_iou, "AS-cDist": aid_cdist, "RS-IoU": 1. - rid_iou, "RS-cDist": rid_cdist } def IoU_cDist( gen_obj_dict, gt_obj_dict, num_states=2, compare_handles=False, iou_include_base=False, rotation_fix_range=True, num_samples=10000, ): """ Compute the IoU-based and centroid-distance-based metrics\n This metric is the average sum of IoU between parts in the two objects over the sampled articulation states and at the resting state\n - gen_obj_dict: the dictionary of the generated object\n - gt_obj_dict: the dictionary of the gt object\n - num_states: the number of articulation states to compute the metric\n - compare_handles (optional): whether to compare the handles\n - iou_include_base (optional): whether to include the base part in the IoU computation\n - rotation_fix_range (optional): whether to fix the rotation range to 90 degrees for revolute joints\n - num_samples (optional): the number of samples to use\n Return:\n - scores: a dictionary of the computed scores\n - "AS-IoU": the average IoU over the articulation states\n - "AS-cDist": the average centroid distance over the articulation states\n - "RS-IoU": the average IoU at the resting state\n - "RS-cDist": the average centroid distance at the resting state\n """ # Make copies of the dictionaries to avoid modifying the original dictionaries gen_dict = deepcopy(gen_obj_dict) gt_dict = deepcopy(gt_obj_dict) # Strip the handles from the object if not comparing them if not compare_handles: gen_dict = remove_handles(gen_dict) gt_dict = remove_handles(gt_dict) # Zero center the objects zero_center_object(gen_dict) zero_center_object(gt_dict) # scale the generated object as a whole to match the size of the gt object gen_bbox_size = compute_overall_bbox_size(gen_dict) gt_bbox_size = compute_overall_bbox_size(gt_dict) scale_factor = gt_bbox_size / gen_bbox_size rescale_object(gen_dict, scale_factor) mapping_gen2gt = find_part_mapping(gen_dict, gt_dict, use_hungarian=True) # for i in range(mapping_gen2gt.shape[0]): # if mapping_gen2gt[i][0] < 100: # gen_dict['diffuse_tree'][i]["parent"] = gt_dict['diffuse_tree'][int(mapping_gen2gt[i][0])]["parent"] # gen_dict['diffuse_tree'][i]["children"] = gt_dict['diffuse_tree'][int(mapping_gen2gt[i][0])]["children"] # gen_dict['diffuse_tree'][i]["id"] = gt_dict['diffuse_tree'][int(mapping_gen2gt[i][0])]["id"] # mapping_gen2gt = find_part_mapping(gen_dict, gt_dict, use_hungarian=True) mapping_gt2gen = find_part_mapping(gt_dict, gen_dict, use_hungarian=True) # Save the original bounding box vertices in rest pose original_gen_bbox_vertices = np.array( [get_bbox_vertices(gen_dict, i) for i in range(len(gen_dict["diffuse_tree"]))], dtype=np.float32, ) original_gt_bbox_vertices = np.array( [get_bbox_vertices(gt_dict, i) for i in range(len(gt_dict["diffuse_tree"]))], dtype=np.float32, ) # import ipdb # ipdb.set_trace() scores_gen2gt = _get_scores( gen_dict, gt_dict, original_gen_bbox_vertices, original_gt_bbox_vertices, mapping_gen2gt, num_states, rotation_fix_range, num_samples, iou_include_base, ) scores_gt2gen = _get_scores( gt_dict, gen_dict, original_gt_bbox_vertices, original_gen_bbox_vertices, mapping_gt2gen, num_states, rotation_fix_range, num_samples, iou_include_base, ) scores = { "AS-IoU": (scores_gen2gt["AS-IoU"] + scores_gt2gen["AS-IoU"]) / 2, "AS-cDist": (scores_gen2gt["AS-cDist"] + scores_gt2gen["AS-cDist"]) / 2, "RS-IoU": (scores_gen2gt["RS-IoU"] + scores_gt2gen["RS-IoU"]) / 2, "RS-cDist": (scores_gen2gt["RS-cDist"] + scores_gt2gen["RS-cDist"]) / 2, } return scores