Merge pull request #2052 from SBCV:feature_mask

Author: alalek

modules/videostab/include/opencv2/videostab/frame_source.hpp

@@ -86,6 +86,28 @@ class CV_EXPORTS VideoFileSource : public IFrameSource
     Ptr<IFrameSource> impl;
 };
 
+class MaskFrameSource : public IFrameSource
+{
+public:
+    MaskFrameSource(const Ptr<IFrameSource>& source): impl(source) {};
+
+    virtual void reset() CV_OVERRIDE { impl->reset(); }
+    virtual Mat nextFrame() CV_OVERRIDE {
+        Mat nextFrame = impl->nextFrame();
+        maskCallback_(nextFrame);
+        return nextFrame;
+    }
+
+    void setMaskCallback(std::function<void(Mat&)> MaskCallback)
+    {
+        maskCallback_ = std::bind(MaskCallback, std::placeholders::_1);
+    };
+
+private:
+    Ptr<IFrameSource> impl;
+    std::function<void(Mat&)> maskCallback_;
+};
+
 //! @}
 
 } // namespace videostab

modules/videostab/include/opencv2/videostab/global_motion.hpp

@@ -180,6 +180,12 @@ class CV_EXPORTS ImageMotionEstimatorBase
     virtual void setMotionModel(MotionModel val) { motionModel_ = val; }
     virtual MotionModel motionModel() const { return motionModel_; }
 
+    virtual void setFrameMask(InputArray mask)
+    {
+        if (!mask.empty())
+            CV_Error(Error::StsNotImplemented, "Mask support is not implemented.");
+    }
+
     virtual Mat estimate(const Mat &frame0, const Mat &frame1, bool *ok = 0) = 0;
 
 protected:
@@ -208,6 +214,8 @@ class CV_EXPORTS ToFileMotionWriter : public ImageMotionEstimatorBase
     virtual void setMotionModel(MotionModel val) CV_OVERRIDE { motionEstimator_->setMotionModel(val); }
     virtual MotionModel motionModel() const CV_OVERRIDE { return motionEstimator_->motionModel(); }
 
+    virtual void setFrameMask(InputArray mask) CV_OVERRIDE { motionEstimator_->setFrameMask(mask); }
+
     virtual Mat estimate(const Mat &frame0, const Mat &frame1, bool *ok = 0) CV_OVERRIDE;
 
 private:
@@ -235,6 +243,8 @@ class CV_EXPORTS KeypointBasedMotionEstimator : public ImageMotionEstimatorBase
     void setOutlierRejector(Ptr<IOutlierRejector> val) { outlierRejector_ = val; }
     Ptr<IOutlierRejector> outlierRejector() const { return outlierRejector_; }
 
+    virtual void setFrameMask(InputArray mask) CV_OVERRIDE { mask_ = mask.getMat(); }
+
     virtual Mat estimate(const Mat &frame0, const Mat &frame1, bool *ok = 0) CV_OVERRIDE;
     Mat estimate(InputArray frame0, InputArray frame1, bool *ok = 0);
 
@@ -243,6 +253,7 @@ class CV_EXPORTS KeypointBasedMotionEstimator : public ImageMotionEstimatorBase
     Ptr<FeatureDetector> detector_;
     Ptr<ISparseOptFlowEstimator> optFlowEstimator_;
     Ptr<IOutlierRejector> outlierRejector_;
+    Mat mask_;
 
     std::vector<uchar> status_;
     std::vector<KeyPoint> keypointsPrev_;
@@ -263,6 +274,8 @@ class CV_EXPORTS KeypointBasedMotionEstimatorGpu : public ImageMotionEstimatorBa
     void setOutlierRejector(Ptr<IOutlierRejector> val) { outlierRejector_ = val; }
     Ptr<IOutlierRejector> outlierRejector() const { return outlierRejector_; }
 
+    virtual void setFrameMask(InputArray mask) CV_OVERRIDE { mask_ = mask.getMat(); }
+
     virtual Mat estimate(const Mat &frame0, const Mat &frame1, bool *ok = 0) CV_OVERRIDE;
     Mat estimate(const cuda::GpuMat &frame0, const cuda::GpuMat &frame1, bool *ok = 0);
 
@@ -271,6 +284,7 @@ class CV_EXPORTS KeypointBasedMotionEstimatorGpu : public ImageMotionEstimatorBa
     Ptr<cuda::CornersDetector> detector_;
     SparsePyrLkOptFlowEstimatorGpu optFlowEstimator_;
     Ptr<IOutlierRejector> outlierRejector_;
+    GpuMat mask_;
 
     cuda::GpuMat frame0_, grayFrame0_, frame1_;
     cuda::GpuMat pointsPrev_, points_;

modules/videostab/include/opencv2/videostab/stabilizer.hpp

@@ -77,6 +77,9 @@ class CV_EXPORTS StabilizerBase
     void setFrameSource(Ptr<IFrameSource> val) { frameSource_ = val; }
     Ptr<IFrameSource> frameSource() const { return frameSource_; }
 
+    void setMaskSource(const Ptr<IFrameSource>& val) { maskSource_ = val; }
+    Ptr<IFrameSource> maskSource() const { return maskSource_; }
+
     void setMotionEstimator(Ptr<ImageMotionEstimatorBase> val) { motionEstimator_ = val; }
     Ptr<ImageMotionEstimatorBase> motionEstimator() const { return motionEstimator_; }
 
@@ -110,6 +113,7 @@ class CV_EXPORTS StabilizerBase
 
     Ptr<ILog> log_;
     Ptr<IFrameSource> frameSource_;
+    Ptr<IFrameSource> maskSource_;
     Ptr<ImageMotionEstimatorBase> motionEstimator_;
     Ptr<DeblurerBase> deblurer_;
     Ptr<InpainterBase> inpainter_;

modules/videostab/samples/videostab.cpp

@@ -89,6 +89,8 @@ void printHelp()
             "      Number of keypoints to find in each frame. The default is 1000.\n"
             "  --local-outlier-rejection=(yes|no)\n"
             "      Perform local outlier rejection. The default is no.\n\n"
+            "  --feature-masks=(file_path|no)\n"
+            "      Load masks from file. The default is no.\n\n"
             "  -sm=, --save-motions=(<file_path>|no)\n"
             "      Save estimated motions into file. The default is no.\n"
             "  -lm=, --load-motions=(<file_path>|no)\n"
@@ -297,6 +299,7 @@ int main(int argc, const char **argv)
                 "{  nkps                    | 1000 | }"
                 "{  extra-kps               | 0 | }"
                 "{  local-outlier-rejection | no | }"
+                "{  feature-masks           | no | }"
                 "{ sm  save-motions         | no | }"
                 "{ lm  load-motions         | no | }"
                 "{ r  radius                | 15 | }"
@@ -461,6 +464,19 @@ int main(int argc, const char **argv)
         stabilizer->setFrameSource(source);
         stabilizer->setMotionEstimator(motionEstBuilder->build());
 
+        if (arg("feature-masks") != "no")
+        {
+            Ptr<MaskFrameSource> maskSource = makePtr<MaskFrameSource>(
+                makePtr<VideoFileSource>(arg("feature-masks")));
+            std::function<void(Mat&)> maskCallback = [](Mat & inputFrame)
+            {
+                cv::cvtColor(inputFrame, inputFrame, cv::COLOR_BGR2GRAY);
+                threshold(inputFrame, inputFrame, 127, 255, THRESH_BINARY);
+            };
+            maskSource->setMaskCallback(maskCallback);
+            stabilizer->setMaskSource(maskSource);
+        }
+
         // cast stabilizer to simple frame source interface to read stabilized frames
         stabilizedFrames.reset(dynamic_cast<IFrameSource*>(stabilizer));
 

modules/videostab/src/global_motion.cpp

@@ -725,7 +725,7 @@ Mat KeypointBasedMotionEstimator::estimate(const Mat &frame0, const Mat &frame1,
 Mat KeypointBasedMotionEstimator::estimate(InputArray frame0, InputArray frame1, bool *ok)
 {
     // find keypoints
-    detector_->detect(frame0, keypointsPrev_);
+    detector_->detect(frame0, keypointsPrev_, mask_);
     if (keypointsPrev_.empty())
         return Mat::eye(3, 3, CV_32F);
 
@@ -815,7 +815,7 @@ Mat KeypointBasedMotionEstimatorGpu::estimate(const cuda::GpuMat &frame0, const
     }
 
     // find keypoints
-    detector_->detect(grayFrame0, pointsPrev_);
+    detector_->detect(grayFrame0, pointsPrev_, mask_);
 
     // find correspondences
     optFlowEstimator_.run(frame0, frame1, pointsPrev_, points_, status_);

modules/videostab/src/stabilizer.cpp

@@ -391,6 +391,9 @@ void TwoPassStabilizer::runPrePassIfNecessary()
         {
             if (frameCount_ > 0)
             {
+                if (maskSource_)
+                    motionEstimator_->setFrameMask(maskSource_->nextFrame());
+
                 motions_.push_back(motionEstimator_->estimate(prevFrame, frame, &ok));
 
                 if (doWobbleSuppression_)