|
- import mindspore.numpy as np
- import mindspore
- import mindspore.nn as nn
- import numpy
- from mindspore import Tensor
- from train.provider import compute_box3d_iou
- from models.model_util import NUM_HEADING_BIN, NUM_SIZE_CLUSTER, NUM_OBJECT_POINT,parse_output_to_tensors,g_mean_size_arr,point_cloud_masking
- from mindspore import ms_function
- mindspore_abs = mindspore.ops.Abs()
- mindspore_minimum = mindspore.ops.Minimum()
- mindspore_reduce_mean = mindspore.ops.ReduceMean(keep_dims=True)
- mindspore_sparse_softmax_cross = mindspore.ops.SparseSoftmaxCrossEntropyWithLogits()
- mindspore_expand_dims = mindspore.ops.ExpandDims()
- mindspore_cos = mindspore.ops.Cos()
- mindspore_sin = mindspore.ops.Sin()
- mindspore_ones = mindspore.ops.Ones()
- mindspore_zeros = mindspore.ops.Zeros()
- mindspore_stack = mindspore.ops.Stack(axis=1)
- mindspore_tile = mindspore.ops.Tile()
- mindspore_transpose = mindspore.ops.Transpose()
- class BatchNorm1d_cus(nn.Cell):
- def __init__(self,num_features):
- super(BatchNorm1d_cus, self).__init__()
- self.bn = nn.BatchNorm2d(num_features)
- self.squeeze = mindspore.ops.Squeeze(3)
- @ms_function
- def construct(self,x):
- x = self.squeeze(self.bn(mindspore_expand_dims(x,3)))
- return x
-
- class get_instance_seg_v1_net(nn.Cell):
- def __init__(self):
- super(get_instance_seg_v1_net, self).__init__()
- self.conv1 = nn.Conv1d(4, 64, 1)
- self.conv2 = nn.Conv1d(64, 64, 1)
- self.conv3 = nn.Conv1d(64, 64, 1)
- self.conv4 = nn.Conv1d(64, 128, 1)
- self.conv5 = nn.Conv1d(128, 1024, 1)
-
- self.conv6 = nn.Conv1d(1091, 512, 1)
- self.conv7 = nn.Conv1d(512, 256, 1)
- self.conv8 = nn.Conv1d(256, 128, 1)
- self.conv9 = nn.Conv1d(128, 128, 1)
- self.conv10 = nn.Conv1d(128, 2, 1)
-
- self.bn1 = BatchNorm1d_cus(64)
- self.bn2 = BatchNorm1d_cus(64)
- self.bn3 = BatchNorm1d_cus(64)
- self.bn4 = BatchNorm1d_cus(128)
- self.bn5 = BatchNorm1d_cus(1024)
- self.bn6 = BatchNorm1d_cus(512)
- self.bn7 = BatchNorm1d_cus(256)
- self.bn8 = BatchNorm1d_cus(128)
- self.bn9 = BatchNorm1d_cus(128)
- self.Concat = mindspore.ops.Concat(1)
- self.relu = nn.ReLU()
- self.dropout=nn.Dropout(keep_prob=0.5)
- @ms_function
- def construct(self,point_cloud, one_hot_vec):
- x = self.relu(self.bn1(self.conv1(point_cloud)))
- x = self.relu(self.bn2(self.conv2(x)))
- point_feat = self.relu(self.bn3(self.conv3(x)))
- x = self.relu(self.bn4(self.conv4(point_feat)))
-
- _,global_feat = mindspore.ops.ArgMaxWithValue(axis=2)(self.relu(self.bn5(self.conv5(x))))
-
- global_feat = self.Concat((global_feat, one_hot_vec))
- global_feat = mindspore_expand_dims(global_feat,2)
- global_feat_expand=mindspore_tile(global_feat,(1,1,1024))
- concat_feat = self.Concat((point_feat, global_feat_expand))
-
- x = self.relu(self.bn6(self.conv6(concat_feat)))
- x = self.relu(self.bn7(self.conv7(x)))
- x = self.relu(self.bn8(self.conv8(x)))
- x = self.relu(self.bn9(self.conv9(x)))
- x = self.dropout(x)
- logits = self.conv10(x)
- return logits
-
- class get_center_regression_net(nn.Cell):
- def __init__(self):
- super(get_center_regression_net, self).__init__()
- self.conv1 = nn.Conv1d(3, 128, 1)
- self.conv2 = nn.Conv1d(128, 128, 1)
- self.conv3 = nn.Conv1d(128, 256, 1)
- self.fc1 = nn.Dense(259,256)
- self.fc2 = nn.Dense(256,128)
- self.fc3 = nn.Dense(128,3)
- self.bn1 = BatchNorm1d_cus(128)
- self.bn2 = BatchNorm1d_cus(128)
- self.bn3 = BatchNorm1d_cus(256)
-
- self.bn4 = nn.BatchNorm1d(256)
- self.bn5 = nn.BatchNorm1d(128)
- self.relu = nn.ReLU()
- self.squeeze = mindspore.ops.Squeeze()
- self.Concat = mindspore.ops.Concat(1)
- @ms_function
- def construct(self,object_point_cloud, one_hot_vec):
- x = self.relu(self.bn1(self.conv1(object_point_cloud)))
- x = self.relu(self.bn2(self.conv2(x)))
- _,x = mindspore.ops.ArgMaxWithValue(axis=2)(self.relu(self.bn3(self.conv3(x))))
-
-
- x = self.Concat((self.squeeze(x),one_hot_vec))
- x = self.relu(self.bn4(self.fc1(x)))
- x = self.relu(self.bn5(self.fc2(x)))
- x = self.fc3(x)
- return x
-
- class get_3d_box_estimation_v1_net(nn.Cell):
- def __init__(self):
- super(get_3d_box_estimation_v1_net, self).__init__()
- self.conv1 = nn.Conv1d(3, 128, 1)
- self.conv2 = nn.Conv1d(128, 128, 1)
- self.conv3 = nn.Conv1d(128, 256, 1)
- self.conv4 = nn.Conv1d(256, 512, 1)
- self.bn1 = BatchNorm1d_cus(128)
- self.bn2 = BatchNorm1d_cus(128)
- self.bn3 = BatchNorm1d_cus(256)
- self.bn4 = BatchNorm1d_cus(512)
- self.fc1 = nn.Dense(515,512)
- self.fc2 = nn.Dense(512,256)
- self.fc3 = nn.Dense(256,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4)
-
- self.bn5 = nn.BatchNorm1d(512)
- self.bn6 = nn.BatchNorm1d(256)
- self.relu = nn.ReLU()
- @ms_function
- def construct(self,object_point_cloud, one_hot_vec):
- x = self.relu(self.bn1(self.conv1(object_point_cloud)))
- x = self.relu(self.bn2(self.conv2(x)))
- x = self.relu(self.bn3(self.conv3(x)))
- x = self.relu(self.bn4(self.conv4(x))).max(2)
-
- concat = mindspore.ops.Concat(1)
- x = concat((x,one_hot_vec))
- x = self.relu(self.bn5(self.fc1(x)))
- x = self.relu(self.bn6(self.fc2(x)))
- x = self.fc3(x)
- return x
-
- class get_modelv1(nn.Cell):
- def __init__(self):
- super(get_modelv1, self).__init__()
- self.instance_seg_v1_net=get_instance_seg_v1_net()
- self.center_regression_net=get_center_regression_net()
- self._3d_box_estimation_v1_net=get_3d_box_estimation_v1_net()
- self.transpose = mindspore.ops.Transpose()
- def construct(self,point_cloud, one_hot_vec):
- end_points = {}
- # 3D Instance Segmentation PointNet
-
- logits=self.instance_seg_v1_net(point_cloud, one_hot_vec)
- end_points['mask_logits'] = self.transpose(logits,(0,2,1))
-
- # Masking
- # select masked points and translate to masked points' centroid
- object_point_cloud_xyz, mask_xyz_mean, mask=point_cloud_masking(point_cloud, logits)
- end_points['mask'] = mask
-
- # T-Net and coordinate translation
- center_delta=self.center_regression_net(object_point_cloud_xyz, one_hot_vec)
- stage1_center = center_delta + mask_xyz_mean # Bx3
- end_points['stage1_center'] = stage1_center
-
- # Get object point cloud in object coordinate
- object_point_cloud_xyz_new = object_point_cloud_xyz - mindspore_expand_dims(center_delta,2)
-
- # Amodel Box Estimation PointNet
- output=self._3d_box_estimation_v1_net(object_point_cloud_xyz_new,one_hot_vec)
-
- # Parse output to 3D box parameters
- end_points=parse_output_to_tensors(output, end_points)
- end_points['center'] = end_points['center_boxnet'] + stage1_center
- #for k,v in end_points.items():
- # print(k,v.shape)
- #print(end_points['center'])
- return end_points
- class Myeval(nn.Metric):
- def __init__(self):
- super(Myeval, self).__init__()
- self.clear()
-
- def clear(self):
- self.correct_obj = 0
- self.obj_num = 0
- self.correct_point = 0
- self.point_num = 0
-
-
- @nn.rearrange_inputs
- def update(self, *inputs):
-
- self.correct_obj = inputs[0]
- self.obj_num = inputs[1]
- self.correct_point = inputs[2]
- self.point_num = inputs[3]
-
- def eval(self):
- print("Segmentation Accuracy: "+str(100*float(self.correct_point)/self.point_num))
- print("Box Estimation Accuracy: "+str(100*float(self.correct_obj)/self.obj_num))
- print()
- return 100*float(self.correct_obj)/self.obj_num
-
- class CustomWithEvalCell(nn.Cell):
- def __init__(self, network):
- super(CustomWithEvalCell, self).__init__(auto_prefix=False)
- self.network = network
- self.total_correct = 0
- self.total_seen = 0
- self.total_seen_class = [0 for _ in range(2)]
- self.total_correct_class = [0 for _ in range(2)]
- self.loss_sum = 0
- self.car_num=0
- self.iou3d_correct_car_cnt = 0
- self.pedestrian_num=0
- self.iou3d_correct_pedestrian_cnt = 0
- self.cyclist_num=0
- self.iou3d_correct_cyclist_cnt = 0
-
- def construct(self, point_cloud, one_hot_vec,mask_label, center_label, heading_class_label, heading_residual_label, size_class_label, size_residual_label):
- end_points = self.network(point_cloud, one_hot_vec)
- iou2ds,iou3ds=compute_box3d_iou(end_points['center'].asnumpy(),end_points['heading_scores'].asnumpy(), end_points['heading_residuals'].asnumpy(),
- end_points['size_scores'].asnumpy(), end_points['size_residuals'].asnumpy(),center_label,heading_class_label.asnumpy(),
- heading_residual_label.asnumpy(),size_class_label.asnumpy(),size_residual_label.asnumpy())
-
- sclass=size_class_label.asnumpy()
- batch_label=mask_label.asnumpy()
- ############使用GPU时注释这两行############
- #iou2ds=np.array(iou2ds)
- #iou3ds=np.array(iou3ds)
- ############使用GPU时注释这两行############
- preds_val = mindspore.ops.Argmax(axis=2)(end_points['mask_logits']).asnumpy()
- correct = (preds_val == mask_label.asnumpy()).sum()
- self.total_correct = correct
- self.total_seen = (len(point_cloud.asnumpy())*1024)
- for l in range(2):
- self.total_seen_class[l] += (batch_label==l).sum()
- self.total_correct_class[l] += Tensor(((preds_val==l) & (batch_label==l)).sum())
- self.car_num = (sclass==0).sum()
- self.pedestrian_num = (sclass==3).sum()
- self.cyclist_num = (sclass==5).sum()
-
- self.iou3d_correct_car_cnt = (Tensor(iou3ds>=0.7).astype(mindspore.float32) * Tensor(sclass==0).astype(mindspore.float32).T).sum()
- self.iou3d_correct_pedestrian_cnt = (Tensor(iou3ds>=0.7).astype(mindspore.float32) * Tensor(sclass==3).astype(mindspore.float32).T).sum()
- self.iou3d_correct_cyclist_cnt = (Tensor(iou3ds>=0.7).astype(mindspore.float32) * Tensor(sclass==5).astype(mindspore.float32).T).sum()
- #print("car:",self.iou3d_correct_car_cnt,"/",self.car_num," pedestrian:",self.iou3d_correct_pedestrian_cnt,"/",self.pedestrian_num," cyclist:",self.iou3d_correct_cyclist_cnt,"/",self.cyclist_num)
-
- return end_points['size_scores'],self.iou3d_correct_car_cnt+self.iou3d_correct_pedestrian_cnt+self.iou3d_correct_cyclist_cnt,self.car_num+self.pedestrian_num+self.cyclist_num,self.total_correct,self.total_seen
-
- class CustomWithLossCell(nn.Cell):
- def __init__(self, backbone, loss_fn):
- super(CustomWithLossCell, self).__init__(auto_prefix=False)
- self._backbone = backbone
- self._loss_fn = loss_fn
- #self.total_num=73508
- #self.num_count=0
- #self.i=0
- #self.mloss=0
-
- def predict(self, point_cloud, one_hot_vec):
- end_points = self._backbone(point_cloud, one_hot_vec)
- return end_points
-
- def construct(self, point_cloud, one_hot_vec,mask_label, center_label, heading_class_label, heading_residual_label, size_class_label, size_residual_label):
- end_points = self._backbone(point_cloud, one_hot_vec)
- #loss=self._loss_fn(mask_label, center_label, heading_class_label, heading_residual_label, size_class_label, size_residual_label,end_points)
- #self.mloss = (self.mloss * self.i + float(loss.asnumpy())) / (self.i + 1)
- #self.i+= 1
- #self.num_count+=point_cloud.shape[0]
- #if self.total_num-self.num_count<point_cloud.shape[0]:
- # print("Loss: %4f"%self.mloss)
- # self.i=0
- # self.mloss=0
- # self.num_count=0
- return self._loss_fn(mask_label, center_label, heading_class_label, heading_residual_label, size_class_label, size_residual_label,end_points)
- @ms_function
- def huber_loss(error, delta):
- abs_error = mindspore_abs(error)
- quadratic = mindspore_minimum(abs_error, delta)
- linear = (abs_error - quadratic)
- losses = 0.5 * quadratic**2 + delta * linear
- return mindspore_reduce_mean(losses)
- @ms_function
- def get_box3d_corners_helper(centers, headings, sizes):
- """ TF layer. Input: (N,3), (N,), (N,3), Output: (N,8,3) """
-
- #print '-----', centers
- N = centers.shape[0]
- slice_op = mindspore.ops.Slice()
- l = slice_op(sizes, [0,0], [-1,1]) # (N,1)
- w = slice_op(sizes, [0,1], [-1,1]) # (N,1)
- h = slice_op(sizes, [0,2], [-1,1]) # (N,1)
- #print l,w,h
- concat = mindspore.ops.Concat(1)
- x_corners = concat([l/2,l/2,-l/2,-l/2,l/2,l/2,-l/2,-l/2]) # (N,8)
- y_corners = concat([h/2,h/2,h/2,h/2,-h/2,-h/2,-h/2,-h/2]) # (N,8)
- z_corners = concat([w/2,-w/2,-w/2,w/2,w/2,-w/2,-w/2,w/2]) # (N,8)
- corners = concat([mindspore_expand_dims(x_corners,1), mindspore_expand_dims(y_corners,1), mindspore_expand_dims(z_corners,1)]) # (N,3,8)
- #print x_corners, y_corners, z_corners
- c = mindspore_cos(headings)
- s = mindspore_sin(headings)
- ones = mindspore_ones((N,), mindspore.float32)
- zeros = mindspore_zeros((N,), mindspore.float32)
- row1 = mindspore_stack([c,zeros,s]) # (N,3)
- row2 = mindspore_stack([zeros,ones,zeros])
- row3 = mindspore_stack([-s,zeros,c])
-
- R = concat([mindspore_expand_dims(row1,1), mindspore_expand_dims(row2,1), mindspore_expand_dims(row3,1)]) # (N,3,3)
- #print row1, row2, row3, R, N
- corners_3d = mindspore.ops.matmul(R, corners) # (N,3,8)
- corners_3d += mindspore_tile(mindspore_expand_dims(centers,2), (1,1,8)) # (N,3,8)
- corners_3d = mindspore_transpose(corners_3d,(0,2,1)) # (N,8,3)
- return corners_3d
-
- def get_box3d_corners(center, heading_residuals, size_residuals):
- """ TF layer.
- Inputs:
- center: (B,3)
- heading_residuals: (B,NH)
- size_residuals: (B,NS,3)
- Outputs:
- box3d_corners: (B,NH,NS,8,3) tensor
- """
- batch_size = center.shape[0]
- heading_bin_centers = mindspore.numpy.arange(0,2*np.pi,2*np.pi/NUM_HEADING_BIN) # (NH,)
- headings = heading_residuals + mindspore_expand_dims(heading_bin_centers, 0) # (B,NH)
-
- mean_sizes = mindspore_expand_dims(Tensor(g_mean_size_arr).astype(mindspore.float32), 0) + size_residuals.astype(mindspore.float32) # (B,NS,1)
- sizes = mean_sizes + size_residuals # (B,NS,3)
- sizes = mindspore_tile(mindspore_expand_dims(sizes,1), (1,NUM_HEADING_BIN,1,1)) # (B,NH,NS,3)
- headings = mindspore_tile(mindspore_expand_dims(headings,-1), (1,1,NUM_SIZE_CLUSTER)) # (B,NH,NS)
- centers = mindspore_tile(mindspore_expand_dims(mindspore_expand_dims(center,1),1), (1,NUM_HEADING_BIN, NUM_SIZE_CLUSTER,1)) # (B,NH,NS,3)
-
- N = batch_size*NUM_HEADING_BIN*NUM_SIZE_CLUSTER
- corners_3d = get_box3d_corners_helper(centers.reshape((N,3)), headings.reshape((N,)), sizes.reshape((N,3)) )
- return corners_3d.reshape( [batch_size, NUM_HEADING_BIN, NUM_SIZE_CLUSTER, 8, 3])
-
- class get_loss(nn.Cell):
- def __init__(self):
- super(get_loss, self).__init__()
- def construct(self, mask_label, center_label, heading_class_label, heading_residual_label, size_class_label, size_residual_label,end_points):
- corner_loss_weight=10.0
- box_loss_weight=1.0
- batch_size=mask_label.shape[0]
- # 3D Segmentation loss
-
- mask_loss = 0
- mask_loss=mindspore_sparse_softmax_cross(end_points['mask_logits'].reshape((batch_size*1024,2)), mask_label.astype(mindspore.int32).reshape((batch_size*1024)))
- # Center regression losses
- center_dist = mindspore.numpy.norm(center_label - end_points['center'],axis=-1)
- center_loss = huber_loss(center_dist, 2.0)
- stage1_center_dist = mindspore.numpy.norm(center_label - end_points['stage1_center'],axis=-1)
- stage1_center_loss = huber_loss(stage1_center_dist, 1.0)
-
- # Heading loss
- heading_class_loss = mindspore_reduce_mean(mindspore_sparse_softmax_cross(end_points['heading_scores'], heading_class_label.astype(mindspore.int32)))
- onehot=nn.OneHot(depth=NUM_HEADING_BIN,on_value=1, off_value=0, axis=-1)
- hcls_onehot = onehot(heading_class_label.astype(mindspore.int32)) # BxNUM_HEADING_BIN
- heading_residual_normalized_label =(heading_residual_label / (np.pi/NUM_HEADING_BIN)).astype(mindspore.float32)
- heading_residual_normalized_loss = huber_loss((end_points['heading_residuals_normalized']*hcls_onehot).sum(1) - heading_residual_normalized_label, 1.0).astype(mindspore.float32)
- # Size loss
- size_class_loss = mindspore_reduce_mean(mindspore_sparse_softmax_cross(end_points['size_scores'], size_class_label.astype(mindspore.int32)))
- onehot=nn.OneHot(depth=NUM_SIZE_CLUSTER,on_value=1, off_value=0, axis=-1)
- scls_onehot = onehot(size_class_label.astype(mindspore.int32)) # BxNUM_SIZE_CLUSTER
- scls_onehot_tiled = mindspore_tile(mindspore_expand_dims(scls_onehot, -1), (1,1,3)) # BxNUM_SIZE_CLUSTERx3
- predicted_size_residual_normalized = (end_points['size_residuals_normalized']*scls_onehot_tiled).sum(1) # Bx3
- mean_size_arr_expand = mindspore_expand_dims(Tensor(g_mean_size_arr,mindspore.float32),0) # 1xNUM_SIZE_CLUSTERx3
- mean_size_label = (scls_onehot_tiled * mean_size_arr_expand).sum(1) # Bx3
- size_residual_label_normalized = size_residual_label / mean_size_label
- size_normalized_dist = mindspore.numpy.norm(size_residual_label_normalized - predicted_size_residual_normalized,axis=-1)
- size_residual_normalized_loss = huber_loss(size_normalized_dist, 1.0)
- # Corner loss
- # We select the predicted corners corresponding to the GT heading bin and size cluster.
- corners_3d = get_box3d_corners(end_points['center'],end_points['heading_residuals'],end_points['size_residuals']) # (B,NH,NS,8,3)
- gt_mask = mindspore_tile(mindspore_expand_dims(hcls_onehot, 2), (1,1,NUM_SIZE_CLUSTER)) *mindspore_tile(mindspore_expand_dims(scls_onehot,1), (1,NUM_HEADING_BIN,1)) # (B,NH,NS)
- corners_3d_pred = (mindspore_expand_dims(mindspore_expand_dims(gt_mask,-1),-1) * corners_3d).sum(axis=(1,2)) # (B,8,3)
- heading_bin_centers = mindspore.numpy.arange(0,2*np.pi,2*np.pi/NUM_HEADING_BIN).astype(mindspore.float32) # (NH,)
- heading_label = mindspore_expand_dims(heading_residual_label,1) + mindspore_expand_dims(heading_bin_centers, 0).astype(mindspore.float32) # (B,NH)
- heading_label = (hcls_onehot*heading_label).sum(1).astype(mindspore.float32)
- mean_sizes = mindspore_expand_dims(Tensor(g_mean_size_arr,mindspore.float32), 0) # (1,NS,3)
- size_label = mean_sizes + mindspore_expand_dims(size_residual_label, 1).astype(mindspore.float32) # (1,NS,3) + (B,1,3) = (B,NS,3)
- size_label = (mindspore_expand_dims(scls_onehot,-1)*size_label).sum(1).astype(mindspore.float32) # (B,3)
- corners_3d_gt = get_box3d_corners_helper(center_label, heading_label, size_label) # (B,8,3)
- corners_3d_gt_flip = get_box3d_corners_helper(center_label, heading_label+np.pi, size_label) # (B,8,3)
- corners_dist = mindspore_minimum(mindspore.numpy.norm(corners_3d_pred - corners_3d_gt,axis=-1),mindspore.numpy.norm(corners_3d_pred - corners_3d_gt_flip,axis=-1))
- corners_loss = huber_loss(corners_dist, 1.0)
- # Weighted sum of all losses
- total_loss = mask_loss + box_loss_weight * (center_loss + heading_class_loss + size_class_loss + heading_residual_normalized_loss*20 + size_residual_normalized_loss*20 + stage1_center_loss + corner_loss_weight*corners_loss)
- return total_loss
-
-
|