Skip to content

Commit

Permalink
fix(trt): apply std to input image correctly
Browse files Browse the repository at this point in the history
  • Loading branch information
Bycob authored and mergify[bot] committed Jun 10, 2021
1 parent 7890401 commit 23bd913
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 19 deletions.
36 changes: 19 additions & 17 deletions src/backends/tensorrt/tensorrtinputconns.cc
Original file line number Diff line number Diff line change
Expand Up @@ -31,28 +31,34 @@ namespace dd
{
cv::Mat converted;
int channels = img.channels();

if (_has_mean_scalar && _mean.size() != size_t(channels))
throw InputConnectorBadParamException(
"mean vector be of size the number of channels ("
+ std::to_string(channels) + ")");

if (!_std.empty() && _std.size() != size_t(channels))
throw InputConnectorBadParamException(
"std vector be of size the number of channels ("
+ std::to_string(channels) + ")");

bool has_std = !_std.empty();
int offset = channels * _height * _width * i;
img.convertTo(converted, CV_32F);
boost::float32_t *fbuf = (boost::float32_t *)(_buf.data());
boost::float32_t *cvbuf = (boost::float32_t *)converted.data;
for (int c = 0; c < channels; ++c)
for (int h = 0; h < _height; ++h)
for (int w = 0; w < _width; ++w)
fbuf[offset++]
= _scale * cvbuf[(converted.cols * h + w) * channels + c];
}

void ImgTensorRTInputFileConn::applyMeanToRTBuf(int channels, int i)
{
int offset = channels * _height * _width * i;
for (int c = 0; c < channels; ++c)
for (int h = 0; h < _height; ++h)
for (int w = 0; w < _width; ++w)
{
int data_index = (c * _height + h) * _width + w;
boost::float32_t *fbuf
= static_cast<boost::float32_t *>(_buf.data());
fbuf[offset + data_index] -= _mean[c];
fbuf[offset]
= _scale * cvbuf[(converted.cols * h + w) * channels + c];
if (_has_mean_scalar)
fbuf[offset] -= _mean[c];
if (has_std)
fbuf[offset] /= _std[c];
++offset;
}
}

Expand Down Expand Up @@ -102,10 +108,6 @@ namespace dd
{
cv::Mat img = this->_images.at(_batch_index);
CVMatToRTBuffer(img, i);
if (_has_mean_scalar)
{
applyMeanToRTBuf(img.channels(), i);
}
}
return i;
}
Expand Down
2 changes: 0 additions & 2 deletions src/backends/tensorrt/tensorrtinputconns.h
Original file line number Diff line number Diff line change
Expand Up @@ -96,8 +96,6 @@ namespace dd
_imgs_size; /**< image sizes, used in detection. */

private:
void applyMeanToRTBuf(int channels, int i);
void applyMeanToRTBuf(float *mean, int channels, int i);
void CVMatToRTBuffer(cv::Mat &img, int i);
};

Expand Down

0 comments on commit 23bd913

Please sign in to comment.