import torch import numpy as np import cv2 from ProposalNetwork.scoring.convex_outline import tracing_outline_robust import ProposalNetwork.utils.spaces as spaces from scipy.spatial import cKDTree from ProposalNetwork.utils.utils import iou_2d, mask_iou, mod_mask_iou def score_point_cloud(point_cloud:torch.Tensor, cubes:list[spaces.Cubes], K:torch.Tensor=None, segmentation_mask:torch.Tensor=None): ''' score the cube according to the density (number of points) of the point cloud in the cube ''' # must normalise the point cloud to have the same density for the entire depth verts = cubes.get_all_corners().squeeze(0) min_x, _, = verts[:,0].min(1); max_x, _ = verts[:,0].max(1) min_y, _, = verts[:,1].min(1); max_y, _ = verts[:,1].max(1) min_z, _, = verts[:,2].min(1); max_z, _ = verts[:,2].max(1) point_cloud_dens = ((point_cloud[:,0].view(-1,1) > min_x) & (point_cloud[:,0].view(-1,1) < max_x) & (point_cloud[:,1].view(-1,1) > min_y) & (point_cloud[:,1].view(-1,1) < max_y) & (point_cloud[:,2].view(-1,1) > min_z) & (point_cloud[:,2].view(-1,1) < max_z)) score = point_cloud_dens.sum(0) # method 1 # just in case this is needed in the future, the function can be found at commit ID: 4a06501c46beda804fd3b8ddfcbb27211f89ef66 # area = cube.get_projected_2d_area(K).item() # if area != 0: # score /= area # method 2 # corners = cube.get_bube_corners(K) # bube_mask = np.zeros(segmentation_mask.shape, dtype=np.uint8) # polygon_points = cv2.convexHull(corners.numpy()) # polygon_points = np.array([polygon_points],dtype=np.int32) # cv2.fillPoly(bube_mask, polygon_points, 1) # normalisation = (bube_mask).sum() # if normalisation != 0: # score = score/normalisation return score def score_iou(gt_box, proposal_box): IoU = iou_2d(gt_box,proposal_box) return IoU def modified_chamfer_distance(set1, set2): tree2 = cKDTree(set2) # For each point in set1 (seg point), find the distance to the nearest point in set2 (bube corner) distances2, _ = tree2.query(set1) return np.mean(distances2) def score_corners(segmentation_mask, bube_corners): mask_np = segmentation_mask.cpu().numpy().astype(np.uint8) # Find contours contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Find the minimum area rectangle around the largest contour if contours: largest_contour = max(contours, key=cv2.contourArea) rect = cv2.minAreaRect(largest_contour) box = cv2.boxPoints(rect) else: # if it fails, set the box as the mean of the bube corners mean_min_x = bube_corners[:,:,0].min(1)[0].mean().cpu().numpy() mean_max_x = bube_corners[:,:,0].max(1)[0].mean().cpu().numpy() mean_min_y = bube_corners[:,:,1].min(1)[0].mean().cpu().numpy() mean_max_y = bube_corners[:,:,1].max(1)[0].mean().cpu().numpy() box = np.array([[mean_min_x, mean_min_y], [mean_max_x, mean_min_y], [mean_max_x, mean_max_y], [mean_min_x, mean_max_y]]) bube_corners = bube_corners.squeeze(0) # remove instance dim scores = torch.zeros(len(bube_corners), device=segmentation_mask.device) for i in range(len(bube_corners)): # Chamfer distance bube corners and box scores[i] = modified_chamfer_distance(box, bube_corners[i].cpu().numpy()) max_score = torch.max(scores) return 1 - scores / max_score def score_segmentation(segmentation_mask, bube_corners): ''' segmentation_mask : Mask bube_corners : List of Lists ''' bube_corners = bube_corners.to(device=segmentation_mask.device) bube_corners = bube_corners.squeeze(0) # remove instance dim scores = torch.zeros(len(bube_corners), device=segmentation_mask.device) for i in range(len(bube_corners)): bube_mask = np.zeros(segmentation_mask.shape, dtype='uint8') # Remove "inner" points (2) and put others in correct order # Calculate the convex hull of the points which also orders points correctly polygon_points = cv2.convexHull(np.array(bube_corners[i])) polygon_points = np.array([polygon_points],dtype=np.int32) cv2.fillPoly(bube_mask, polygon_points, 1) scores[i] = mask_iou(segmentation_mask[::4,::4], bube_mask[::4,::4]) return scores def score_mod_segmentation(segmentation_mask, bube_corners): ''' segmentation_mask : Mask bube_corners : List of Lists ''' bube_corners = bube_corners.to(device=segmentation_mask.device) bube_corners = bube_corners.squeeze(0) # remove instance dim scores = torch.zeros(len(bube_corners), device=segmentation_mask.device) for i in range(len(bube_corners)): bube_mask = np.zeros(segmentation_mask.shape, dtype='uint8') # Remove "inner" points (2) and put others in correct order # Calculate the convex hull of the points which also orders points correctly polygon_points = cv2.convexHull(np.array(bube_corners[i])) polygon_points = np.array([polygon_points],dtype=np.int32) cv2.fillPoly(bube_mask, polygon_points, 1) scores[i] = mod_mask_iou(segmentation_mask[::4,::4], bube_mask[::4,::4]) return scores def score_segmentation_v2(segmentation_mask, pred_cubes, K): scores = [] for i in range(len(pred_cubes.tensor.squeeze())): v_2d = pred_cubes[:, i].get_bube_corners(K).squeeze() _, f = pred_cubes[:, i].get_cuboids_verts_faces() f = f.squeeze() points, ids = tracing_outline_robust(v_2d.numpy(), f.numpy()) # not doing any projection,just simply take the verts's x and y . bube_mask = np.zeros(segmentation_mask.shape, dtype='uint8') # append first point to close the loop # points = np.append(points, [points[0]], axis=0) cv2.fillPoly(bube_mask, np.expand_dims(points,0).astype(int), 1) scores.append(mask_iou(segmentation_mask, bube_mask)) return scores def score_dimensions(category, dimensions, gt_boxes, pred_boxes): ''' category : List dimensions : List of Lists P(dim|priors) ''' # category_name = Metadatacatalog.thing_classes[category] # for printing and checking that correct [prior_mean, prior_std] = category dimensions_scores = torch.exp(-1/2 * ((dimensions - prior_mean)/prior_std)**2) scores = dimensions_scores.mean(1) gt_ratio = (gt_boxes.tensor[0,2]-gt_boxes.tensor[0,0])/(gt_boxes.tensor[0,3]-gt_boxes.tensor[0,1]) pred_ratios = (pred_boxes.tensor[:,2]-pred_boxes.tensor[:,0])/(pred_boxes.tensor[:,3]-pred_boxes.tensor[:,1]) differences = torch.abs(gt_ratio-pred_ratios) max_difference = torch.max(differences) return (1 - differences / max_difference) * scores def score_ratios(gt_box,pred_boxes): gt_points = gt_box.tensor[0] differences = torch.abs(pred_boxes.tensor - gt_points).sum(axis=1) max_difference = torch.max(differences) return 1 - differences / max_difference # 3D Dim Ratio gt_ratio = gt_dim[0]/gt_dim[1] pred_ratios = pred_dims[:,0]/pred_dims[:,1] differences = torch.abs(pred_ratios-gt_ratio) max_difference = torch.max(differences) return 1 - differences / max_difference # 2D Dim Ratio gt_ratio = (gt_dim.tensor[0,2]-gt_dim.tensor[0,0])/(gt_dim.tensor[0,3]-gt_dim.tensor[0,1]) pred_ratios = (pred_dims.tensor[:,2]-pred_dims.tensor[:,0])/(pred_dims.tensor[:,3]-pred_dims.tensor[:,1]) differences = torch.abs(pred_ratios-gt_ratio) max_difference = torch.max(differences) return 1 - differences / max_difference def score_function(gt_box, proposal_box, bube_corners, segmentation_mask, category, dimensions): score = 1.0 score *= score_iou(gt_box, proposal_box) score *= score_segmentation(bube_corners, segmentation_mask) score *= score_dimensions(category, dimensions) return score if __name__ == '__main__': # testing s = score_point_cloud(torch.tensor([[0.1,0.1,0.1],[0.2,0.2,0.2],[-3,0,0]]), [spaces.Cube(torch.tensor([0.5,0.5,0.5,1,1,1]), torch.eye(3))]) print(s) assert s == 2