176 lines
9.3 KiB
Plaintext
176 lines
9.3 KiB
Plaintext
#include "pose_engine.hpp"
|
|
#include "common.hpp"
|
|
#include "gpu_preprocess.hpp"
|
|
|
|
#include <opencv2/imgproc.hpp>
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
|
|
namespace fangjia {
|
|
namespace {
|
|
|
|
constexpr int kInputWidth = 288;
|
|
constexpr int kInputHeight = 384;
|
|
constexpr int kKeypoints = 133;
|
|
constexpr int kSimccX = 576;
|
|
constexpr int kSimccY = 768;
|
|
|
|
__global__ void poseNormalizeKernel(
|
|
float* output, std::uint8_t const* input, int people, int inputStep) {
|
|
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int pixels = kInputWidth * kInputHeight;
|
|
int total = people * pixels;
|
|
if (index >= total) return;
|
|
int person = index / pixels;
|
|
int pixel = index - person * pixels;
|
|
int y = pixel / kInputWidth;
|
|
int x = pixel - y * kInputWidth;
|
|
auto const* value = input + person * inputStep * kInputHeight + y * inputStep + x * 3;
|
|
float* destination = output + static_cast<std::size_t>(person) * 3 * pixels;
|
|
destination[pixel] = static_cast<float>((static_cast<double>(value[0]) - 123.675) / 58.395);
|
|
destination[pixels + pixel] = static_cast<float>((static_cast<double>(value[1]) - 116.28) / 57.12);
|
|
destination[2 * pixels + pixel] = static_cast<float>((static_cast<double>(value[2]) - 103.53) / 57.375);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
RtmwEngine::RtmwEngine(
|
|
std::filesystem::path const& enginePath, TrtLogger& logger,
|
|
int maxSourceWidth, int maxSourceHeight, int maxPeople)
|
|
: engine_(enginePath, logger), maxSourceWidth_(maxSourceWidth),
|
|
maxSourceHeight_(maxSourceHeight), maxPeople_(maxPeople),
|
|
centersHost_(maxPeople * 2), scalesHost_(maxPeople * 2),
|
|
simccXHost_(static_cast<std::size_t>(maxPeople) * kKeypoints * kSimccX),
|
|
simccYHost_(static_cast<std::size_t>(maxPeople) * kKeypoints * kSimccY) {
|
|
engine_.setInputShape("input", nvinfer1::Dims4{1,3,kInputHeight,kInputWidth});
|
|
engine_.prepare();
|
|
auto xs=engine_.shape("simcc_x"),ys=engine_.shape("simcc_y");
|
|
if(xs.nbDims!=3||ys.nbDims!=3||xs.d[0]!=1||ys.d[0]!=1||xs.d[1]!=kKeypoints||ys.d[1]!=kKeypoints||xs.d[2]!=kSimccX||ys.d[2]!=kSimccY)
|
|
throw std::runtime_error("RTMW engine I/O shape does not match the 133-point model");
|
|
for(auto name:{"input","simcc_x","simcc_y"})if(engine_.dataType(name)!=nvinfer1::DataType::kFLOAT)
|
|
throw std::runtime_error("RTMW requires float32 I/O bindings");
|
|
std::size_t const values = static_cast<std::size_t>(maxPeople) * 3 * kInputWidth * kInputHeight;
|
|
CUDA_CHECK(cudaMalloc(
|
|
&cropsDevice_, static_cast<std::size_t>(maxPeople_) * kInputWidth * kInputHeight * 3));
|
|
CUDA_CHECK(cudaMalloc(&networkDevice_, values * sizeof(float)));
|
|
CUDA_CHECK(cudaMalloc(&sourceDevice_, static_cast<size_t>(maxSourceWidth_)*maxSourceHeight_*3));
|
|
}
|
|
|
|
RtmwEngine::~RtmwEngine() {
|
|
if (sourceDevice_) cudaFree(sourceDevice_);
|
|
if (cropsDevice_) cudaFree(cropsDevice_);
|
|
if (networkDevice_) cudaFree(networkDevice_);
|
|
}
|
|
|
|
void RtmwEngine::warmup(cudaStream_t stream) {
|
|
std::vector<Box> boxes{{0, 0, static_cast<float>(maxSourceWidth_), static_cast<float>(maxSourceHeight_), 1, 1}};
|
|
cv::Mat black(maxSourceHeight_, maxSourceWidth_, CV_8UC3, cv::Scalar(0, 0, 0));
|
|
for (int i = 0; i < 2; ++i) infer(black, boxes, stream);
|
|
}
|
|
|
|
std::vector<PersonPose> RtmwEngine::infer(
|
|
cv::Mat const& frame, std::vector<Box> const& boxes, cudaStream_t stream,
|
|
float* preprocessMs, float* inferenceMs, float* postprocessMs) {
|
|
int batch = std::min(static_cast<int>(boxes.size()), maxPeople_);
|
|
if (batch <= 0) return {};
|
|
auto const preprocessStarted = std::chrono::steady_clock::now();
|
|
int const cropStep = kInputWidth * 3;
|
|
int const pixels = kInputWidth * kInputHeight;
|
|
if (frame.cols>maxSourceWidth_ || frame.rows>maxSourceHeight_) throw std::runtime_error("Pose input exceeds allocation");
|
|
auto npp = cms::nppContext(stream);
|
|
CUDA_CHECK(cudaMemcpy2DAsync(sourceDevice_, frame.cols*3, frame.data, frame.step, frame.cols*3, frame.rows, cudaMemcpyHostToDevice, stream));
|
|
for (int i = 0; i < batch; ++i) {
|
|
auto const& box = boxes[i];
|
|
double width = (static_cast<double>(box.x2) - static_cast<double>(box.x1)) * 1.25;
|
|
double height = (static_cast<double>(box.y2) - static_cast<double>(box.y1)) * 1.25;
|
|
constexpr double aspect = static_cast<double>(kInputWidth) / kInputHeight;
|
|
if (width > height * aspect) height = width / aspect;
|
|
else width = height * aspect;
|
|
centersHost_[i * 2] = 0.5 * (static_cast<double>(box.x1) + static_cast<double>(box.x2));
|
|
centersHost_[i * 2 + 1] = 0.5 * (static_cast<double>(box.y1) + static_cast<double>(box.y2));
|
|
scalesHost_[i * 2] = width;
|
|
scalesHost_[i * 2 + 1] = height;
|
|
|
|
cv::Point2f source[3]{
|
|
{static_cast<float>(centersHost_[i * 2]), static_cast<float>(centersHost_[i * 2 + 1])},
|
|
{static_cast<float>(centersHost_[i * 2]), static_cast<float>(centersHost_[i * 2 + 1] - width * 0.5)},
|
|
{static_cast<float>(centersHost_[i * 2] - width * 0.5), static_cast<float>(centersHost_[i * 2 + 1] - width * 0.5)}};
|
|
cv::Point2f destination[3]{
|
|
{kInputWidth * 0.5F, kInputHeight * 0.5F},
|
|
{kInputWidth * 0.5F, kInputHeight * 0.5F - kInputWidth * 0.5F},
|
|
{0.0F, kInputHeight * 0.5F - kInputWidth * 0.5F}};
|
|
cv::Mat affine = cv::getAffineTransform(source, destination);
|
|
auto* cropDevice=cropsDevice_+static_cast<size_t>(i)*cropStep*kInputHeight;
|
|
Npp8u fill[3]={0,0,0};
|
|
cms::nppCheck(nppiSet_8u_C3R_Ctx(fill,cropDevice,cropStep,{kInputWidth,kInputHeight},npp));
|
|
double coeff[2][3];
|
|
for(int row=0;row<2;++row) for(int col=0;col<3;++col) coeff[row][col]=affine.at<double>(row,col);
|
|
cms::nppCheck(nppiWarpAffine_8u_C3R_Ctx(sourceDevice_,{frame.cols,frame.rows},frame.cols*3,{0,0,frame.cols,frame.rows},cropDevice,cropStep,{0,0,kInputWidth,kInputHeight},coeff,NPPI_INTER_LINEAR,npp));
|
|
}
|
|
int const totalPixels = batch * pixels;
|
|
poseNormalizeKernel<<<(totalPixels + 255) / 256, 256, 0, stream>>>(
|
|
networkDevice_, cropsDevice_, batch, cropStep);
|
|
CUDA_CHECK(cudaGetLastError());
|
|
CUDA_CHECK(cudaStreamSynchronize(stream));
|
|
if (preprocessMs) {
|
|
*preprocessMs = static_cast<float>(std::chrono::duration<double, std::milli>(
|
|
std::chrono::steady_clock::now() - preprocessStarted).count());
|
|
}
|
|
|
|
// RTMPose's Python API invokes the ONNX/TensorRT session once per person.
|
|
// Preserve that execution contract instead of batching people together.
|
|
engine_.setInputShape("input", nvinfer1::Dims4{1, 3, kInputHeight, kInputWidth});
|
|
engine_.prepare();
|
|
CudaEventTimer timer;
|
|
float inferTime = 0.0F;
|
|
std::size_t const inputValues = static_cast<std::size_t>(3) * pixels;
|
|
std::size_t const xValues = static_cast<std::size_t>(kKeypoints) * kSimccX;
|
|
std::size_t const yValues = static_cast<std::size_t>(kKeypoints) * kSimccY;
|
|
for (int person = 0; person < batch; ++person) {
|
|
CUDA_CHECK(cudaMemcpyAsync(
|
|
engine_.device("input"), networkDevice_ + static_cast<std::size_t>(person) * inputValues,
|
|
inputValues * sizeof(float), cudaMemcpyDeviceToDevice, stream));
|
|
timer.start(stream);
|
|
if (!engine_.enqueue(stream)) throw std::runtime_error("RTMW enqueue failed");
|
|
inferTime += timer.stop(stream);
|
|
CUDA_CHECK(cudaMemcpyAsync(
|
|
simccXHost_.data() + static_cast<std::size_t>(person) * xValues,
|
|
engine_.device("simcc_x"), xValues * sizeof(float), cudaMemcpyDeviceToHost, stream));
|
|
CUDA_CHECK(cudaMemcpyAsync(
|
|
simccYHost_.data() + static_cast<std::size_t>(person) * yValues,
|
|
engine_.device("simcc_y"), yValues * sizeof(float), cudaMemcpyDeviceToHost, stream));
|
|
}
|
|
if (inferenceMs) *inferenceMs = inferTime;
|
|
|
|
auto const postprocessStarted = std::chrono::steady_clock::now();
|
|
CUDA_CHECK(cudaStreamSynchronize(stream));
|
|
std::vector<PersonPose> result(batch);
|
|
for (int person = 0; person < batch; ++person) {
|
|
result[person].box = boxes[person];
|
|
for (int keypoint = 0; keypoint < kKeypoints; ++keypoint) {
|
|
auto xBegin = simccXHost_.begin() + (person * kKeypoints + keypoint) * kSimccX;
|
|
auto yBegin = simccYHost_.begin() + (person * kKeypoints + keypoint) * kSimccY;
|
|
auto xIt = std::max_element(xBegin, xBegin + kSimccX);
|
|
auto yIt = std::max_element(yBegin, yBegin + kSimccY);
|
|
float score = 0.5F * (*xIt + *yIt);
|
|
if(!std::isfinite(score))throw std::runtime_error("Non-finite RTMW output");
|
|
float localX = score > 0.0F ? static_cast<float>(xIt - xBegin) * 0.5F : -0.5F;
|
|
float localY = score > 0.0F ? static_cast<float>(yIt - yBegin) * 0.5F : -0.5F;
|
|
double scaleX = scalesHost_[person * 2];
|
|
double scaleY = scalesHost_[person * 2 + 1];
|
|
result[person].keypoints[keypoint] = {
|
|
static_cast<float>(static_cast<double>(localX) / kInputWidth * scaleX + centersHost_[person * 2] - scaleX * 0.5),
|
|
static_cast<float>(static_cast<double>(localY) / kInputHeight * scaleY + centersHost_[person * 2 + 1] - scaleY * 0.5),
|
|
score};
|
|
}
|
|
}
|
|
float post = static_cast<float>(std::chrono::duration<double, std::milli>(
|
|
std::chrono::steady_clock::now() - postprocessStarted).count());
|
|
if (postprocessMs) *postprocessMs = post;
|
|
return result;
|
|
}
|
|
|
|
} // namespace fangjia
|