Optimized augmentation for speed

zhengshunjie · Feb 28, 2018 · 6d7e794 · 6d7e794
1 parent 6e7d170
commit 6d7e794
Show file tree

Hide file tree

Showing 4 changed files with 161 additions and 109 deletions.
diff --git a/include/caffe/openpose/dataAugmentation.hpp b/include/caffe/openpose/dataAugmentation.hpp
@@ -24,9 +24,9 @@ namespace caffe {
     // Cropping
     cv::Point2i estimateCrop(const MetaData& metaData, const OPTransformationParameter& param_);
     void applyCrop(cv::Mat& imageAugmented, const cv::Point2i& cropCenter, const cv::Mat& image,
-                   const unsigned char defaultBorderValue, const OPTransformationParameter& param_);
+                   const unsigned char defaultBorderValue, const cv::Size& cropSize);
     void applyCrop(MetaData& metaData, const cv::Point2i& cropCenter,
-                   const OPTransformationParameter& param_, const PoseModel poseModel);
+                   const cv::Size& cropSize, const PoseModel poseModel);
     // Flipping
     bool estimateFlip(const MetaData& metaData,
                       const OPTransformationParameter& param_);

diff --git a/include/caffe/openpose/oPDataTransformer.hpp b/include/caffe/openpose/oPDataTransformer.hpp
@@ -89,7 +89,7 @@ class OPDataTransformer {
     void generateDataAndLabel(Dtype* transformedData, Dtype* transformedLabel, const Datum& datum,
                               const Datum* datumNegative);
     void generateDepthLabelMap(Dtype* transformedLabel, const cv::Mat& depth) const;
-    void generateLabelMap(Dtype* transformedLabel, const cv::Mat& image, const cv::Mat& maskMiss,
+    void generateLabelMap(Dtype* transformedLabel, const cv::Size& imageSize, const cv::Mat& maskMiss,
                           const MetaData& metaData) const;
     void putGaussianMaps(Dtype* entry, const cv::Point2f& center, const int stride, const int gridX, const int gridY,
                          const float sigma) const;

diff --git a/src/caffe/openpose/dataAugmentation.cpp b/src/caffe/openpose/dataAugmentation.cpp
@@ -166,7 +166,7 @@ namespace caffe {
     }
 
     void applyCrop(cv::Mat& imageAugmented, const cv::Point2i& cropCenter, const cv::Mat& image,
-                   const unsigned char defaultBorderValue, const OPTransformationParameter& param_)
+                   const unsigned char defaultBorderValue, const cv::Size& cropSize)
     {
         if (!image.empty())
         {
@@ -175,8 +175,8 @@ namespace caffe {
                 throw std::runtime_error{"Input and output images must be different"
                                          + getLine(__LINE__, __FUNCTION__, __FILE__)};
             // Parameters
-            const auto cropX = (int) param_.crop_size_x();
-            const auto cropY = (int) param_.crop_size_y();
+            const auto cropX = (int)cropSize.width;
+            const auto cropY = (int)cropSize.height;
             // Crop image
             // 1. Allocate memory
             imageAugmented = cv::Mat(cropY, cropX, image.type(), cv::Scalar{(double)defaultBorderValue});
@@ -227,11 +227,11 @@ namespace caffe {
     }
 
     void applyCrop(MetaData& metaData, const cv::Point2i& cropCenter,
-                   const OPTransformationParameter& param_, const PoseModel poseModel)
+                   const cv::Size& cropSize, const PoseModel poseModel)
     {
         // Update metaData
-        const auto cropX = (int) param_.crop_size_x();
-        const auto cropY = (int) param_.crop_size_y();
+        const auto cropX = (int)cropSize.width;
+        const auto cropY = (int)cropSize.height;
         const int offsetLeft = -(cropCenter.x - (cropX/2));
         const int offsetUp = -(cropCenter.y - (cropY/2));
         const cv::Point2f offsetPoint{(float)offsetLeft, (float)offsetUp};