videostab.cpp 14 KB
Newer Older
wester committed
1 2 3 4
#if defined __clang__
#  pragma GCC diagnostic ignored "-Wdelete-non-virtual-dtor"
#endif

wester committed
5 6 7 8 9
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <stdexcept>
wester committed
10 11 12 13 14
#include "opencv2/core/core.hpp"
#include "opencv2/video/video.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/videostab/videostab.hpp"
wester committed
15 16 17 18 19

using namespace std;
using namespace cv;
using namespace cv::videostab;

wester committed
20

wester committed
21 22 23 24 25 26 27 28 29 30
Ptr<IFrameSource> stabilizedFrames;
string saveMotionsPath;
double outputFps;
string outputPath;
bool quietMode;

void run();
void saveMotionsIfNecessary();
void printHelp();

wester committed
31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
class GlobalMotionReader : public IGlobalMotionEstimator
{
public:
    GlobalMotionReader(string path)
    {
        ifstream f(path.c_str());
        if (!f.is_open())
            throw runtime_error("can't open motions file: " + path);
        int size = 0; f >> size;
        motions_.resize(size);
        for (int i = 0; i < size; ++i)
        {
            Mat_<float> M(3, 3);
            for (int l = 0; l < 3; ++l)
                for (int s = 0; s < 3; ++s)
                    f >> M(l,s);
            motions_[i] = M;
        }
        pos_ = 0;
    }

    virtual Mat estimate(const Mat &/*frame0*/, const Mat &/*frame1*/)
    {
        if (pos_ >= motions_.size())
        {
            stringstream text;
            text << "can't load motion between frames " << pos_ << " and " << pos_+1;
            throw runtime_error(text.str());
        }
        return motions_[pos_++];
    }

private:
    vector<Mat> motions_;
    size_t pos_;
};
wester committed
67 68 69 70 71 72 73 74

void run()
{
    VideoWriter writer;
    Mat stabilizedFrame;

    while (!(stabilizedFrame = stabilizedFrames->nextFrame()).empty())
    {
wester committed
75 76
        if (!saveMotionsPath.empty())
            saveMotionsIfNecessary();
wester committed
77 78 79
        if (!outputPath.empty())
        {
            if (!writer.isOpened())
wester committed
80
                writer.open(outputPath, CV_FOURCC('X','V','I','D'), outputFps, stabilizedFrame.size());
wester committed
81 82 83 84 85 86
            writer << stabilizedFrame;
        }
        if (!quietMode)
        {
            imshow("stabilizedFrame", stabilizedFrame);
            char key = static_cast<char>(waitKey(3));
wester committed
87 88
            if (key == 27)
                break;
wester committed
89 90 91
        }
    }

wester committed
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119
    cout << "\nfinished\n";
}


void saveMotionsIfNecessary()
{
    static bool areMotionsSaved = false;
    if (!areMotionsSaved)
    {
        IFrameSource *frameSource = static_cast<IFrameSource*>(stabilizedFrames);
        TwoPassStabilizer *twoPassStabilizer = dynamic_cast<TwoPassStabilizer*>(frameSource);
        if (twoPassStabilizer)
        {
            ofstream f(saveMotionsPath.c_str());
            const vector<Mat> &motions = twoPassStabilizer->motions();
            f << motions.size() << endl;
            for (size_t i = 0; i < motions.size(); ++i)
            {
                Mat_<float> M = motions[i];
                for (int l = 0, k = 0; l < 3; ++l)
                    for (int s = 0; s < 3; ++s, ++k)
                        f << M(l,s) << " ";
                f << endl;
            }
        }
        areMotionsSaved = true;
        cout << "motions are saved";
    }
wester committed
120 121 122 123 124 125 126 127
}


void printHelp()
{
    cout << "OpenCV video stabilizer.\n"
            "Usage: videostab <file_path> [arguments]\n\n"
            "Arguments:\n"
wester committed
128
            "  -m, --model=(transl|transl_and_scale|linear_sim|affine)\n"
wester committed
129 130
            "      Set motion model. The default is affine.\n"
            "  --outlier-ratio=<float_number>\n"
wester committed
131
            "      Outliers ratio in motion estimation. The default is 0.5.\n"
wester committed
132
            "  --min-inlier-ratio=<float_number>\n"
wester committed
133 134 135 136 137 138 139 140 141 142
            "      Minimum inlier ratio to decide if estimated motion is OK. The default is 0.1,\n"
            "      but you may want to increase it.\n\n"
            "  --save-motions=<file_path>\n"
            "      Save estimated motions into file.\n"
            "  --load-motions=<file_path>\n"
            "      Load motions from file.\n\n"
            "  -r, --radius=<int_number>\n"
            "      Set smoothing radius. The default is 15.\n"
            "  --stdev=<float_number>\n"
            "      Set smoothing weights standard deviation. The default is sqrt(radius).\n\n"
wester committed
143 144 145 146
            "  --deblur=(yes|no)\n"
            "      Do deblurring.\n"
            "  --deblur-sens=<float_number>\n"
            "      Set deblurring sensitivity (from 0 to +inf). The default is 0.1.\n\n"
wester committed
147 148 149 150 151 152
            "  -t, --trim-ratio=<float_number>\n"
            "      Set trimming ratio (from 0 to 0.5). The default is 0.0.\n"
            "  --est-trim=(yes|no)\n"
            "      Estimate trim ratio automatically. The default is yes (that leads to two passes,\n"
            "      you can turn it off if you want to use one pass only).\n"
            "  --incl-constr=(yes|no)\n"
wester committed
153
            "      Ensure the inclusion constraint is always satisfied. The default is no.\n\n"
wester committed
154
            "  --border-mode=(replicate|reflect|const)\n"
wester committed
155 156 157 158 159
            "      Set border extrapolation mode. The default is replicate.\n\n"
            "  --mosaic=(yes|no)\n"
            "      Do consistent mosaicing. The default is no.\n"
            "  --mosaic-stdev=<float_number>\n"
            "      Consistent mosaicing stdev threshold. The default is 10.0.\n\n"
wester committed
160 161 162
            "  --motion-inpaint=(yes|no)\n"
            "      Do motion inpainting (requires GPU support). The default is no.\n"
            "  --dist-thresh=<float_number>\n"
wester committed
163
            "      Estimated flow distance threshold for motion inpainting. The default is 5.0.\n\n"
wester committed
164
            "  --color-inpaint=(no|average|ns|telea)\n"
wester committed
165
            "      Do color inpainting. The defailt is no.\n"
wester committed
166 167 168
            "  --color-inpaint-radius=<float_number>\n"
            "      Set color inpainting radius (for ns and telea options only).\n\n"
            "  -o, --output=(no|<file_path>)\n"
wester committed
169
            "      Set output file path explicitely. The default is stabilized.avi.\n"
wester committed
170 171
            "  --fps=<int_number>\n"
            "      Set output video FPS explicitely. By default the source FPS is used.\n"
wester committed
172 173 174
            "  -q, --quiet\n"
            "      Don't show output video frames.\n\n"
            "  -h, --help\n"
wester committed
175 176
            "      Print help.\n"
            "\n";
wester committed
177 178 179 180 181 182 183 184
}


int main(int argc, const char **argv)
{
    try
    {
        const char *keys =
wester committed
185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
                "{ 1 | | | | }"
                "{ m | model | | }"
                "{ | min-inlier-ratio | | }"
                "{ | outlier-ratio | | }"
                "{ | save-motions | | }"
                "{ | load-motions | | }"
                "{ r | radius | | }"
                "{ | stdev | | }"
                "{ | deblur | | }"
                "{ | deblur-sens | | }"
                "{ | est-trim | yes | }"
                "{ t | trim-ratio | | }"
                "{ | incl-constr | | }"
                "{ | border-mode | | }"
                "{ | mosaic | | }"
                "{ | mosaic-stdev | | }"
                "{ | motion-inpaint | | }"
                "{ | dist-thresh | | }"
                "{ | color-inpaint | no | }"
                "{ | color-inpaint-radius | | }"
                "{ o | output | stabilized.avi | }"
                "{ | fps | | }"
                "{ q | quiet | false | }"
                "{ h | help | false | }";
wester committed
209 210 211 212
        CommandLineParser cmd(argc, argv, keys);

        // parse command arguments

wester committed
213
        if (cmd.get<bool>("help"))
wester committed
214 215 216 217 218
        {
            printHelp();
            return 0;
        }

wester committed
219 220 221 222
        StabilizerBase *stabilizer;
        GaussianMotionFilter *motionFilter = 0;

        if (!cmd.get<string>("stdev").empty())
wester committed
223
        {
wester committed
224 225
            motionFilter = new GaussianMotionFilter();
            motionFilter->setStdev(cmd.get<float>("stdev"));
wester committed
226 227
        }

wester committed
228 229
        if (!cmd.get<string>("save-motions").empty())
            saveMotionsPath = cmd.get<string>("save-motions");
wester committed
230 231

        bool isTwoPass =
wester committed
232 233
                cmd.get<string>("est-trim") == "yes" ||
                !cmd.get<string>("save-motions").empty();
wester committed
234 235 236 237

        if (isTwoPass)
        {
            TwoPassStabilizer *twoPassStabilizer = new TwoPassStabilizer();
wester committed
238 239 240 241
            if (!cmd.get<string>("est-trim").empty())
                twoPassStabilizer->setEstimateTrimRatio(cmd.get<string>("est-trim") == "yes");
            if (motionFilter)
                twoPassStabilizer->setMotionStabilizer(motionFilter);
wester committed
242 243 244 245
            stabilizer = twoPassStabilizer;
        }
        else
        {
wester committed
246 247 248
            OnePassStabilizer *onePassStabilizer= new OnePassStabilizer();
            if (motionFilter)
                onePassStabilizer->setMotionFilter(motionFilter);
wester committed
249 250 251
            stabilizer = onePassStabilizer;
        }

wester committed
252 253 254
        string inputPath = cmd.get<string>("1");
        if (inputPath.empty())
            throw runtime_error("specify video file path");
wester committed
255

wester committed
256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272
        VideoFileSource *frameSource = new VideoFileSource(inputPath);
        outputFps = frameSource->fps();
        stabilizer->setFrameSource(frameSource);
        cout << "frame count: " << frameSource->frameCount() << endl;

        PyrLkRobustMotionEstimator *motionEstimator = new PyrLkRobustMotionEstimator();
        if (cmd.get<string>("model") == "transl")
            motionEstimator->setMotionModel(TRANSLATION);
        else if (cmd.get<string>("model") == "transl_and_scale")
            motionEstimator->setMotionModel(TRANSLATION_AND_SCALE);
        else if (cmd.get<string>("model") == "linear_sim")
            motionEstimator->setMotionModel(LINEAR_SIMILARITY);
        else if (cmd.get<string>("model") == "affine")
            motionEstimator->setMotionModel(AFFINE);
        else if (!cmd.get<string>("model").empty())
            throw runtime_error("unknow motion mode: " + cmd.get<string>("model"));
        if (!cmd.get<string>("outlier-ratio").empty())
wester committed
273
        {
wester committed
274 275 276
            RansacParams ransacParams = motionEstimator->ransacParams();
            ransacParams.eps = cmd.get<float>("outlier-ratio");
            motionEstimator->setRansacParams(ransacParams);
wester committed
277
        }
wester committed
278 279 280 281 282
        if (!cmd.get<string>("min-inlier-ratio").empty())
            motionEstimator->setMinInlierRatio(cmd.get<float>("min-inlier-ratio"));
        stabilizer->setMotionEstimator(motionEstimator);
        if (!cmd.get<string>("load-motions").empty())
            stabilizer->setMotionEstimator(new GlobalMotionReader(cmd.get<string>("load-motions")));
wester committed
283

wester committed
284 285
        if (!cmd.get<string>("radius").empty())
            stabilizer->setRadius(cmd.get<int>("radius"));
wester committed
286

wester committed
287
        if (cmd.get<string>("deblur") == "yes")
wester committed
288
        {
wester committed
289 290 291
            WeightingDeblurer *deblurer = new WeightingDeblurer();
            if (!cmd.get<string>("deblur-sens").empty())
                deblurer->setSensitivity(cmd.get<float>("deblur-sens"));
wester committed
292 293 294
            stabilizer->setDeblurer(deblurer);
        }

wester committed
295 296 297 298 299
        if (!cmd.get<string>("trim-ratio").empty())
            stabilizer->setTrimRatio(cmd.get<float>("trim-ratio"));

        if (!cmd.get<string>("incl-constr").empty())
            stabilizer->setCorrectionForInclusion(cmd.get<string>("incl-constr") == "yes");
wester committed
300

wester committed
301
        if (cmd.get<string>("border-mode") == "reflect")
wester committed
302
            stabilizer->setBorderMode(BORDER_REFLECT);
wester committed
303
        else if (cmd.get<string>("border-mode") == "replicate")
wester committed
304
            stabilizer->setBorderMode(BORDER_REPLICATE);
wester committed
305
        else if (cmd.get<string>("border-mode") == "const")
wester committed
306
            stabilizer->setBorderMode(BORDER_CONSTANT);
wester committed
307 308
        else if (!cmd.get<string>("border-mode").empty())
            throw runtime_error("unknown border extrapolation mode: " + cmd.get<string>("border-mode"));
wester committed
309 310

        InpaintingPipeline *inpainters = new InpaintingPipeline();
wester committed
311
        if (cmd.get<string>("mosaic") == "yes")
wester committed
312
        {
wester committed
313 314 315 316
            ConsistentMosaicInpainter *inpainter = new ConsistentMosaicInpainter();
            if (!cmd.get<string>("mosaic-stdev").empty())
                inpainter->setStdevThresh(cmd.get<float>("mosaic-stdev"));
            inpainters->pushBack(inpainter);
wester committed
317
        }
wester committed
318
        if (cmd.get<string>("motion-inpaint") == "yes")
wester committed
319
        {
wester committed
320 321 322 323
            MotionInpainter *inpainter = new MotionInpainter();
            if (!cmd.get<string>("dist-thresh").empty())
                inpainter->setDistThreshold(cmd.get<float>("dist-thresh"));
            inpainters->pushBack(inpainter);
wester committed
324
        }
wester committed
325
        if (!cmd.get<string>("color-inpaint").empty())
wester committed
326
        {
wester committed
327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347
            if (cmd.get<string>("color-inpaint") == "average")
                inpainters->pushBack(new ColorAverageInpainter());
            else if (!cmd.get<string>("color-inpaint-radius").empty())
            {
                float radius = cmd.get<float>("color-inpaint-radius");
                if (cmd.get<string>("color-inpaint") == "ns")
                    inpainters->pushBack(new ColorInpainter(INPAINT_NS, radius));
                else if (cmd.get<string>("color-inpaint") == "telea")
                    inpainters->pushBack(new ColorInpainter(INPAINT_TELEA, radius));
                else if (cmd.get<string>("color-inpaint") != "no")
                    throw runtime_error("unknown color inpainting method: " + cmd.get<string>("color-inpaint"));
            }
            else
            {
                if (cmd.get<string>("color-inpaint") == "ns")
                    inpainters->pushBack(new ColorInpainter(INPAINT_NS));
                else if (cmd.get<string>("color-inpaint") == "telea")
                    inpainters->pushBack(new ColorInpainter(INPAINT_TELEA));
                else if (cmd.get<string>("color-inpaint") != "no")
                    throw runtime_error("unknown color inpainting method: " + cmd.get<string>("color-inpaint"));
            }
wester committed
348
        }
wester committed
349 350
        if (!inpainters->empty())
            stabilizer->setInpainter(inpainters);
wester committed
351

wester committed
352
        stabilizer->setLog(new LogToStdout());
wester committed
353

wester committed
354 355 356 357 358 359 360 361
        outputPath = cmd.get<string>("output") != "no" ? cmd.get<string>("output") : "";

        if (!cmd.get<string>("fps").empty())
            outputFps = cmd.get<double>("fps");

        quietMode = cmd.get<bool>("quiet");

        stabilizedFrames = dynamic_cast<IFrameSource*>(stabilizer);
wester committed
362 363 364 365 366 367 368 369 370 371 372 373

        run();
    }
    catch (const exception &e)
    {
        cout << "error: " << e.what() << endl;
        stabilizedFrames.release();
        return -1;
    }
    stabilizedFrames.release();
    return 0;
}