|
- import mindspore
- import mindspore.numpy as np
- from mindspore import dtype as mstype
- import numpy
- from mindspore import Tensor
- from mindspore import ms_function
-
- NUM_HEADING_BIN = 12
- NUM_SIZE_CLUSTER = 8 # one cluster for each type
- NUM_OBJECT_POINT = 512
- g_type2class={'Car':0, 'Van':1, 'Truck':2, 'Pedestrian':3,
- 'Person_sitting':4, 'Cyclist':5, 'Tram':6, 'Misc':7}
- g_class2type = {g_type2class[t]:t for t in g_type2class}
- g_type2onehotclass = {'Car': 0, 'Pedestrian': 1, 'Cyclist': 2}
- g_type_mean_size = {'Car': numpy.array([3.88311640418,1.62856739989,1.52563191462]),
- 'Van': numpy.array([5.06763659,1.9007158,2.20532825]),
- 'Truck': numpy.array([10.13586957,2.58549199,3.2520595]),
- 'Pedestrian': numpy.array([0.84422524,0.66068622,1.76255119]),
- 'Person_sitting': numpy.array([0.80057803,0.5983815,1.27450867]),
- 'Cyclist': numpy.array([1.76282397,0.59706367,1.73698127]),
- 'Tram': numpy.array([16.17150617,2.53246914,3.53079012]),
- 'Misc': numpy.array([3.64300781,1.54298177,1.92320313])}
- g_mean_size_arr = numpy.zeros((NUM_SIZE_CLUSTER, 3)) # size clustrs
- for i in range(NUM_SIZE_CLUSTER):
- g_mean_size_arr[i,:] = g_type_mean_size[g_class2type[i]]
- mindspore_expand_dims = mindspore.ops.ExpandDims()
- mindspore_transpose = mindspore.ops.Transpose()
- mindspore_tile = mindspore.ops.Tile()
- mindspore_squeeze2 = mindspore.ops.Squeeze(2)
- mindspore_squeeze1 = mindspore.ops.Squeeze(1)
- mindspore_maximum = mindspore.ops.Maximum()
- mindspore_slice = mindspore.ops.Slice()
- mindspore_concat = mindspore.ops.Concat(-1)
- mindspore_GatherNd = mindspore.ops.GatherNd()
- mindspore_reshape = mindspore.ops.Reshape()
- g_mean_size_arr=Tensor(g_mean_size_arr)
-
- def parse_output_to_tensors(output, end_points):
- ''' Parse batch output to separate tensors (added to end_points)
- Input:
- output: TF tensor in shape (B,3+2*NUM_HEADING_BIN+4*NUM_SIZE_CLUSTER)
- end_points: dict
- Output:
- end_points: dict (updated)
- '''
- batch_size = output.shape[0]
- center = mindspore_slice(output, (0,0), (-1,3))
- end_points['center_boxnet'] = center
-
- heading_scores = mindspore_slice(output, (0,3), (-1,NUM_HEADING_BIN))
- heading_residuals_normalized = mindspore_slice(output, (0,3+NUM_HEADING_BIN),
- (-1,NUM_HEADING_BIN))
- end_points['heading_scores'] = heading_scores # BxNUM_HEADING_BIN
- end_points['heading_residuals_normalized'] = heading_residuals_normalized # BxNUM_HEADING_BIN (-1 to 1)
- end_points['heading_residuals'] = heading_residuals_normalized * (np.pi/NUM_HEADING_BIN) # BxNUM_HEADING_BIN
-
- size_scores = mindspore_slice(output, (0,3+NUM_HEADING_BIN*2),
- (-1,NUM_SIZE_CLUSTER)) # BxNUM_SIZE_CLUSTER
- size_residuals_normalized = mindspore_slice(output,
- (0,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER), (-1,NUM_SIZE_CLUSTER*3))
- size_residuals_normalized = mindspore_reshape(size_residuals_normalized,(batch_size, NUM_SIZE_CLUSTER, 3)) # BxNUM_SIZE_CLUSTERx3
- end_points['size_scores'] = size_scores
- end_points['size_residuals_normalized'] = size_residuals_normalized
- end_points['size_residuals'] = size_residuals_normalized *mindspore_expand_dims(Tensor(g_mean_size_arr, dtype=mstype.float32), 0)
-
- return end_points
-
- def tf_gather_object_pc(point_cloud, mask, npoints=512):
- ''' Gather object point clouds according to predicted masks.
- Input:
- point_cloud: TF tensor in shape (B,N,C)
- mask: TF tensor in shape (B,N) of 0 (not pick) or 1 (pick)
- npoints: int scalar, maximum number of points to keep (default: 512)
- Output:
- object_pc: TF tensor in shape (B,npoint,C)
- indices: TF int tensor in shape (B,npoint,2)
- '''
- def mask_to_indices(mask):
- indices = numpy.zeros((mask.shape[0], npoints, 2), dtype=numpy.int32)
- for i in range(mask.shape[0]):
- pos_indices = numpy.where(mask[i,:]>0.5)[0]
- # skip cases when pos_indices is empty
- if len(pos_indices) > 0:
- if len(pos_indices) > npoints:
- choice = numpy.random.choice(len(pos_indices),
- npoints, replace=False)
- else:
- choice = numpy.random.choice(len(pos_indices),
- npoints-len(pos_indices), replace=True)
- choice = numpy.concatenate((numpy.arange(len(pos_indices)), choice))
- numpy.random.shuffle(choice)
- indices[i,:,1] = pos_indices[choice]
- indices[i,:,0] = i
- return Tensor(indices)
- mask=mask.asnumpy()
- indices = mask_to_indices(mask)
- #print(indices.shape)
- object_pc = mindspore_GatherNd(point_cloud, indices)
- return object_pc, indices
-
- def point_cloud_masking(point_cloud, logits, xyz_only=True):
- ''' Select point cloud with predicted 3D mask,
- translate coordinates to the masked points centroid.
-
- Input:
- point_cloud: TF tensor in shape (B,N,C)
- logits: TF tensor in shape (B,N,2)
- end_points: dict
- xyz_only: boolean, if True only return XYZ channels
- Output:
- object_point_cloud: TF tensor in shape (B,M,3)
- for simplicity we only keep XYZ here
- M = NUM_OBJECT_POINT as a hyper-parameter
- mask_xyz_mean: TF tensor in shape (B,3)
- '''
- point_cloud=mindspore_transpose(point_cloud,(0,2,1))
- logits=mindspore_transpose(logits,(0,2,1))
- batch_size = point_cloud.shape[0]
- num_point = point_cloud.shape[1]
- mask = mindspore_slice(logits,(0,0,0),(-1,-1,1)) < mindspore_slice(logits,(0,0,1),(-1,-1,1))
- mask = mask.astype("float32")
- mask_count = mindspore_tile(mask.sum(axis=1,keepdims=True),(1,1,3)) # Bx1x3
- point_cloud_xyz = mindspore_slice(point_cloud, (0,0,0), (-1,-1,3)) # BxNx3
- mask_xyz_mean = (mindspore_tile(mask, (1,1,3))*point_cloud_xyz).sum(axis=1, keepdims=True) # Bx1x3
-
- mask = mindspore_squeeze2(mask) # BxN
-
- mask_xyz_mean = mask_xyz_mean/mindspore_maximum(mask_count,1) # Bx1x3
-
- # Translate to masked points' centroid
- point_cloud_xyz_stage1 = point_cloud_xyz - \
- mindspore_tile(mask_xyz_mean, (1,num_point,1))
-
- if xyz_only:
- point_cloud_stage1 = point_cloud_xyz_stage1
- else:
- point_cloud_features = mindspore_slice(point_cloud, (0,0,3), (-1,-1,-1))
- point_cloud_stage1 = mindspore_concat((point_cloud_xyz_stage1, point_cloud_features))
- num_channels = point_cloud_stage1.shape[2]
-
- object_point_cloud, _ = tf_gather_object_pc(point_cloud_stage1,mask, NUM_OBJECT_POINT)
-
- object_point_cloud= mindspore_reshape(object_point_cloud,(batch_size, NUM_OBJECT_POINT, num_channels))
-
- return mindspore_transpose(object_point_cloud,(0,2,1)), mindspore_squeeze1(mask_xyz_mean), mask
|