wehub-resource-sync
189 lines
7.2 KiB
Plaintext
189 lines
7.2 KiB
Plaintext
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <iostream>
|
|
#include <cuda_runtime_api.h>
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
#define checkCudaErrors(status_) \
|
|
{ \
|
|
auto const status = status_; \
|
|
if (status != 0) \
|
|
{ \
|
|
std::cout << "Cuda failure: " << cudaGetErrorString(status) \
|
|
<< " at line " << __LINE__ \
|
|
<< " in file " << __FILE__ \
|
|
<< " error status: " << status \
|
|
<< std::endl; \
|
|
exit(EXIT_FAILURE); \
|
|
} \
|
|
}
|
|
|
|
#ifndef M_PI
|
|
#define M_PI 3.14159265358979323846
|
|
#endif
|
|
|
|
__device__ float sigmoid(const float x) { return 1.0f / (1.0f + expf(-x)); }
|
|
|
|
__global__ void postprocess_kernal(const float *cls_input,
|
|
float const* box_input,
|
|
const float *dir_cls_input,
|
|
float *anchors,
|
|
float *anchors_bottom_height,
|
|
float *bndbox_output,
|
|
int *object_counter,
|
|
const float min_x_range,
|
|
const float max_x_range,
|
|
const float min_y_range,
|
|
const float max_y_range,
|
|
const int feature_x_size,
|
|
const int feature_y_size,
|
|
const int num_anchors,
|
|
const int num_classes,
|
|
const int num_box_values,
|
|
const float score_thresh,
|
|
const float dir_offset,
|
|
const float dir_limit_offset,
|
|
const int num_dir_bins)
|
|
{
|
|
int max_box_num = feature_x_size * feature_y_size * num_anchors;
|
|
int loc_index =blockIdx.x;
|
|
int batch_idx = blockIdx.x / (feature_x_size * feature_y_size);
|
|
int loc_index_in_frame = blockIdx.x % (feature_x_size * feature_y_size);
|
|
int ith_anchor = threadIdx.x;
|
|
if (ith_anchor >= num_anchors)
|
|
{
|
|
return;
|
|
}
|
|
int col = loc_index_in_frame % feature_x_size;
|
|
int row = loc_index_in_frame / feature_x_size;
|
|
float x_offset = min_x_range + col * (max_x_range - min_x_range) / (feature_x_size - 1);
|
|
float y_offset = min_y_range + row * (max_y_range - min_y_range) / (feature_y_size - 1);
|
|
int cls_offset = loc_index * num_classes * num_anchors + ith_anchor * num_classes;
|
|
float dev_cls[2] = {-1, 0};
|
|
const float *scores = cls_input + cls_offset;
|
|
float max_score = sigmoid(scores[0]);
|
|
int cls_id = 0;
|
|
for (int i = 1; i < num_classes; i++) {
|
|
float cls_score = sigmoid(scores[i]);
|
|
if (cls_score > max_score) {
|
|
max_score = cls_score;
|
|
cls_id = i;
|
|
}
|
|
}
|
|
dev_cls[0] = static_cast<float>(cls_id);
|
|
dev_cls[1] = max_score;
|
|
if (dev_cls[1] >= score_thresh)
|
|
{
|
|
int box_offset = loc_index * num_anchors * num_box_values + ith_anchor * num_box_values;
|
|
int dir_cls_offset = loc_index * num_anchors * 2 + ith_anchor * 2;
|
|
float *anchor_ptr = anchors + ith_anchor * 4;
|
|
float z_offset = anchor_ptr[2] / 2 + anchors_bottom_height[ith_anchor / 2];
|
|
float anchor[7] = {x_offset, y_offset, z_offset, anchor_ptr[0], anchor_ptr[1], anchor_ptr[2], anchor_ptr[3]};
|
|
float const* box_encodings = box_input + box_offset;
|
|
float xa = anchor[0];
|
|
float ya = anchor[1];
|
|
float za = anchor[2];
|
|
float dxa = anchor[3];
|
|
float dya = anchor[4];
|
|
float dza = anchor[5];
|
|
float ra = anchor[6];
|
|
float diagonal = sqrtf(dxa * dxa + dya * dya);
|
|
float be0 = box_encodings[0] * diagonal + xa;
|
|
float be1 = box_encodings[1] * diagonal + ya;
|
|
float be2 = box_encodings[2] * dza + za;
|
|
float be3 = expf(box_encodings[3]) * dxa;
|
|
float be4 = expf(box_encodings[4]) * dya;
|
|
float be5 = expf(box_encodings[5]) * dza;
|
|
float be6 = box_encodings[6] + ra;
|
|
float yaw;
|
|
int dir_label = dir_cls_input[dir_cls_offset] > dir_cls_input[dir_cls_offset + 1] ? 0 : 1;
|
|
float period = 2.0f * float(M_PI) / num_dir_bins;
|
|
float val = be6 - dir_offset;
|
|
float dir_rot = val - floor(val / period + dir_limit_offset) * period;
|
|
yaw = dir_rot + dir_offset + period * dir_label;
|
|
int resCount = atomicAdd(object_counter + batch_idx, 1);
|
|
float *data = bndbox_output + (batch_idx * max_box_num + resCount) * 9;
|
|
data[0] = be0;
|
|
data[1] = be1;
|
|
data[2] = be2;
|
|
data[3] = be3;
|
|
data[4] = be4;
|
|
data[5] = be5;
|
|
data[6] = yaw;
|
|
data[7] = dev_cls[0];
|
|
data[8] = dev_cls[1];
|
|
}
|
|
}
|
|
|
|
|
|
void decodeBbox3DLaunch(
|
|
const int batch_size,
|
|
const float *cls_input,
|
|
const float *box_input,
|
|
const float *dir_cls_input,
|
|
float *anchors,
|
|
float *anchors_bottom_height,
|
|
float *bndbox_output,
|
|
int *object_counter,
|
|
const float min_x_range,
|
|
const float max_x_range,
|
|
const float min_y_range,
|
|
const float max_y_range,
|
|
const int feature_x_size,
|
|
const int feature_y_size,
|
|
const int num_anchors,
|
|
const int num_classes,
|
|
const int num_box_values,
|
|
const float score_thresh,
|
|
const float dir_offset,
|
|
const float dir_limit_offset,
|
|
const int num_dir_bins,
|
|
cudaStream_t stream)
|
|
{
|
|
int bev_size = batch_size * feature_x_size * feature_y_size;
|
|
dim3 threads (num_anchors);
|
|
dim3 blocks (bev_size);
|
|
postprocess_kernal<<<blocks, threads, 0, stream>>>
|
|
(cls_input,
|
|
box_input,
|
|
dir_cls_input,
|
|
anchors,
|
|
anchors_bottom_height,
|
|
bndbox_output,
|
|
object_counter,
|
|
min_x_range,
|
|
max_x_range,
|
|
min_y_range,
|
|
max_y_range,
|
|
feature_x_size,
|
|
feature_y_size,
|
|
num_anchors,
|
|
num_classes,
|
|
num_box_values,
|
|
score_thresh,
|
|
dir_offset,
|
|
dir_limit_offset,
|
|
num_dir_bins);
|
|
checkCudaErrors(cudaGetLastError());
|
|
}
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|