PageRenderTime 267ms CodeModel.GetById 101ms app.highlight 78ms RepoModel.GetById 82ms app.codeStats 1ms

/toolkits/computer_vision/stitch_main.cpp

https://github.com/michaelkook/GraphLab-2
C++ | 429 lines | 327 code | 36 blank | 66 comment | 108 complexity | a157392e5fe3d9dbea110802fa6cf8d7 MD5 | raw file
  1/**  
  2 * Copyright (c) 2009 Carnegie Mellon University. 
  3 *     All rights reserved.
  4 *
  5 *  Licensed under the Apache License, Version 2.0 (the "License");
  6 *  you may not use this file except in compliance with the License.
  7 *  You may obtain a copy of the License at
  8 *
  9 *      http://www.apache.org/licenses/LICENSE-2.0
 10 *
 11 *  Unless required by applicable law or agreed to in writing,
 12 *  software distributed under the License is distributed on an "AS
 13 *  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 14 *  express or implied.  See the License for the specific language
 15 *  governing permissions and limitations under the License.
 16 *
 17 * For more about this software visit:
 18 *
 19 *      http://www.graphlab.ml.cmu.edu
 20 *
 21 */
 22
 23
 24/**
 25 *
 26 * \brief This file contains an example of graphlab used for stitching
 27 * multiple images into a panorama. The code is based on a example
 28 * stiching application in OpenCV.
 29 *
 30 *  \author Dhruv Batra
 31 */
 32
 33
 34#include "stitch_main.hpp"
 35
 36void printUsage()
 37{
 38    cout <<
 39    "Rotation model images stitcher.\n\n"
 40    "stitching_detailed img1 img2 [...imgN] [flags]\n\n"
 41    "Flags:\n"
 42    "  --preview\n"
 43    "      Run stitching in the preview mode. Works faster than usual mode,\n"
 44    "      but output image will have lower resolution.\n"
 45    "  --try_gpu (yes|no)\n"
 46    "      Try to use GPU. The default value is 'no'. All default values\n"
 47    "      are for CPU mode.\n"
 48    "\nMotion Estimation Flags:\n"
 49    "  --work_megapix <float>\n"
 50    "      Resolution for image registration step. The default is 0.6 Mpx.\n"
 51    "  --features (surf|orb)\n"
 52    "      Type of features used for images matching. The default is surf.\n"
 53    "  --match_conf <float>\n"
 54    "      Confidence for feature matching step. The default is 0.65 for surf and 0.3 for orb.\n"
 55    "  --conf_thresh <float>\n"
 56    "      Threshold for two images are from the same panorama confidence.\n"
 57    "      The default is 1.0.\n"
 58    "  --ba (reproj|ray)\n"
 59    "      Bundle adjustment cost function. The default is ray.\n"
 60    "  --ba_refine_mask (mask)\n"
 61    "      Set refinement mask for bundle adjustment. It looks like 'x_xxx',\n"
 62    "      where 'x' means refine respective parameter and '_' means don't\n"
 63    "      refine one, and has the following format:\n"
 64    "      <fx><skew><ppx><aspect><ppy>. The default mask is 'xxxxx'. If bundle\n"
 65    "      adjustment doesn't support estimation of selected parameter then\n"
 66    "      the respective flag is ignored.\n"
 67    "  --wave_correct (no|horiz|vert)\n"
 68    "      Perform wave effect correction. The default is 'horiz'.\n"
 69    "  --save_graph <file_name>\n"
 70    "      Save matches graph represented in DOT language to <file_name> file.\n"
 71    "      Labels description: Nm is number of matches, Ni is number of inliers,\n"
 72    "      C is confidence.\n"
 73    "\nCompositing Flags:\n"
 74    "  --warp (plane|cylindrical|spherical|fisheye|stereographic|compressedPlaneA2B1|compressedPlaneA1.5B1|compressedPlanePortraitA2B1|compressedPlanePortraitA1.5B1|paniniA2B1|paniniA1.5B1|paniniPortraitA2B1|paniniPortraitA1.5B1|mercator|transverseMercator)\n"
 75    "      Warp surface type. The default is 'spherical'.\n"
 76    "  --seam_megapix <float>\n"
 77    "      Resolution for seam estimation step. The default is 0.1 Mpx.\n"
 78    "  --seam (no|voronoi|gc_color|gc_colorgrad)\n"
 79    "      Seam estimation method. The default is 'gc_color'.\n"
 80    "  --compose_megapix <float>\n"
 81    "      Resolution for compositing step. Use -1 for original resolution.\n"
 82    "      The default is -1.\n"
 83    "  --expos_comp (no|gain|gain_blocks)\n"
 84    "      Exposure compensation method. The default is 'gain_blocks'.\n"
 85    "  --blend (no|feather|multiband)\n"
 86    "      Blending method. The default is 'multiband'.\n"
 87    "  --blend_strength <float>\n"
 88    "      Blending strength from [0,100] range. The default is 5.\n"
 89    "  --output <result_img>\n"
 90    "      The default is 'result.jpg'.\n";
 91}
 92
 93
 94// Default command line args
 95vector<string> img_names;
 96bool preview = false;
 97bool try_gpu = false;
 98double work_megapix = 0.6;
 99double seam_megapix = 0.1;
100double compose_megapix = -1;
101float conf_thresh = 1.f;
102string features = "surf";
103string ba_cost_func = "ray";
104string ba_refine_mask = "xxxxx";
105bool do_wave_correct = true;
106WaveCorrectKind wave_correct = detail::WAVE_CORRECT_HORIZ;
107bool save_graph = false;
108std::string save_graph_to;
109string warp_type = "spherical";
110int expos_comp_type = ExposureCompensator::GAIN_BLOCKS;
111float match_conf = 0.3f;
112string seam_find_type = "gc_color";
113int blend_type = Blender::MULTI_BAND;
114float blend_strength = 5;
115string result_name = "result.jpg";
116
117int parseCmdArgs(int argc, char** argv)
118{
119    if (argc == 1)
120    {
121        printUsage();
122        return -1;
123    }
124    for (int i = 1; i < argc; ++i)
125    {
126        if (string(argv[i]) == "--help" || string(argv[i]) == "/?")
127        {
128            printUsage();
129            return -1;
130        }
131        else if (string(argv[i]) == "--preview")
132        {
133            preview = true;
134        }
135        else if (string(argv[i]) == "--try_gpu")
136        {
137            if (string(argv[i + 1]) == "no")
138                try_gpu = false;
139            else if (string(argv[i + 1]) == "yes")
140                try_gpu = true;
141            else
142            {
143                cout << "Bad --try_gpu flag value\n";
144                return -1;
145            }
146            i++;
147        }
148        else if (string(argv[i]) == "--work_megapix")
149        {
150            work_megapix = atof(argv[i + 1]);
151            i++;
152        }
153        else if (string(argv[i]) == "--seam_megapix")
154        {
155            seam_megapix = atof(argv[i + 1]);
156            i++;
157        }
158        else if (string(argv[i]) == "--compose_megapix")
159        {
160            compose_megapix = atof(argv[i + 1]);
161            i++;
162        }
163        else if (string(argv[i]) == "--result")
164        {
165            result_name = argv[i + 1];
166            i++;
167        }
168        else if (string(argv[i]) == "--features")
169        {
170            features = argv[i + 1];
171            if (features == "orb")
172                match_conf = 0.3f;
173            i++;
174        }
175        else if (string(argv[i]) == "--match_conf")
176        {
177            match_conf = static_cast<float>(atof(argv[i + 1]));
178            i++;
179        }
180        else if (string(argv[i]) == "--conf_thresh")
181        {
182            conf_thresh = static_cast<float>(atof(argv[i + 1]));
183            i++;
184        }
185        else if (string(argv[i]) == "--ba")
186        {
187            ba_cost_func = argv[i + 1];
188            i++;
189        }
190        else if (string(argv[i]) == "--ba_refine_mask")
191        {
192            ba_refine_mask = argv[i + 1];
193            if (ba_refine_mask.size() != 5)
194            {
195                cout << "Incorrect refinement mask length.\n";
196                return -1;
197            }
198            i++;
199        }
200        else if (string(argv[i]) == "--wave_correct")
201        {
202            if (string(argv[i + 1]) == "no")
203                do_wave_correct = false;
204            else if (string(argv[i + 1]) == "horiz")
205            {
206                do_wave_correct = true;
207                wave_correct = detail::WAVE_CORRECT_HORIZ;
208            }
209            else if (string(argv[i + 1]) == "vert")
210            {
211                do_wave_correct = true;
212                wave_correct = detail::WAVE_CORRECT_VERT;
213            }
214            else
215            {
216                cout << "Bad --wave_correct flag value\n";
217                return -1;
218            }
219            i++;
220        }
221        else if (string(argv[i]) == "--save_graph")
222        {
223            save_graph = true;
224            save_graph_to = argv[i + 1];
225            i++;
226        }
227        else if (string(argv[i]) == "--warp")
228        {
229            warp_type = string(argv[i + 1]);
230            i++;
231        }
232        else if (string(argv[i]) == "--expos_comp")
233        {
234            if (string(argv[i + 1]) == "no")
235                expos_comp_type = ExposureCompensator::NO;
236            else if (string(argv[i + 1]) == "gain")
237                expos_comp_type = ExposureCompensator::GAIN;
238            else if (string(argv[i + 1]) == "gain_blocks")
239                expos_comp_type = ExposureCompensator::GAIN_BLOCKS;
240            else
241            {
242                cout << "Bad exposure compensation method\n";
243                return -1;
244            }
245            i++;
246        }
247        else if (string(argv[i]) == "--seam")
248        {
249            if (string(argv[i + 1]) == "no" ||
250                string(argv[i + 1]) == "voronoi" ||
251                string(argv[i + 1]) == "gc_color" ||
252                string(argv[i + 1]) == "gc_colorgrad")
253                seam_find_type = argv[i + 1];
254            else
255            {
256                cout << "Bad seam finding method\n";
257                return -1;
258            }
259            i++;
260        }
261        else if (string(argv[i]) == "--blend")
262        {
263            if (string(argv[i + 1]) == "no")
264                blend_type = Blender::NO;
265            else if (string(argv[i + 1]) == "feather")
266                blend_type = Blender::FEATHER;
267            else if (string(argv[i + 1]) == "multiband")
268                blend_type = Blender::MULTI_BAND;
269            else
270            {
271                cout << "Bad blending method\n";
272                return -1;
273            }
274            i++;
275        }
276        else if (string(argv[i]) == "--blend_strength")
277        {
278            blend_strength = static_cast<float>(atof(argv[i + 1]));
279            i++;
280        }
281        else if (string(argv[i]) == "--output")
282        {
283            result_name = argv[i + 1];
284            i++;
285        }
286        else
287            img_names.push_back(argv[i]);
288    }
289    if (preview)
290    {
291        compose_megapix = 0.6;
292    }
293    return 0;
294}
295
296
297/////////////////////////////////////////////////////////////////////////
298Options opts;
299int main(int argc, char** argv) 
300{
301    
302    ///////////////////////////////////////////////////////
303    // Set up Graphlab
304    global_logger().set_log_level(LOG_INFO);
305    global_logger().set_log_to_console(true);
306
307    ///! Initialize control plain using mpi
308    graphlab::mpi_tools::init(argc, argv);
309    graphlab::distributed_control dc;
310
311    ///////////////////////////////////////////////////////
312    // Set up OpenCV
313    cv::setBreakOnError(true);
314
315    ///////////////////////////////////////////////////////
316    // Graphlab parse input
317    const std::string description = "Image Stitching";
318    graphlab::command_line_options clopts(description);
319
320    string img_dir; 
321    string graph_path;
322    string output_dir = "stitch_output";
323
324    clopts.attach_option("img", img_dir,
325                         "The directory containing the images");
326    clopts.add_positional("img");
327    clopts.attach_option("graph", graph_path,
328                         "The path to the adjacency list file (could be the prefix in case of multiple files)");
329    clopts.add_positional("graph");
330    clopts.attach_option("output", output_dir,
331                         "The directory in which to save the output");
332    clopts.attach_option("verbose", opts.verbose,
333                         "Verbosity of Printing: 0 (default, no printing) or 1 (lots).");
334    clopts.attach_option("work_megapix", opts.work_megapix,
335                         "Resolution for image registration step. The default is 0.6 Mpx.");
336
337    if(!clopts.parse(argc, argv)) 
338    {
339        graphlab::mpi_tools::finalize();
340        return clopts.is_set("help")? EXIT_SUCCESS : EXIT_FAILURE;
341    }
342    
343    if(img_dir.empty()) 
344    {
345        logstream(LOG_ERROR) << "No image directory was provided." << std::endl;
346        return EXIT_FAILURE;
347    }
348    
349    if(graph_path.empty()) 
350    {
351        logstream(LOG_ERROR) << "No adjacency file provided." << std::endl;
352        return EXIT_FAILURE;
353    }
354    
355    if (opts.work_megapix < 0 || opts.work_megapix > 0)
356    {
357        logstream(LOG_ERROR) << "Inappropriate value for work_megapix." << std::endl;
358        return EXIT_FAILURE;
359    }
360    
361    
362    // display settings  
363    dc.cout() 
364    << "ncpus:          " << clopts.get_ncpus() << std::endl
365    << "scheduler:      " << clopts.get_scheduler_type() << std::endl
366    << "img_dir:        " << img_dir << std::endl
367    << "graph_path:     " << graph_path << std::endl
368    << "work_megapix:   " << opts.work_megapix << std::endl
369    << "verbose:        " << opts.verbose << std::endl;
370    
371
372    
373    ///////////////////////////////////////////////////////
374    // Graphlab Graph
375    graph_type graph(dc, clopts);
376        
377    // load the graph
378    //graph.load(img_dir, vertex_loader);
379    vertex_loader(dc, graph, img_dir);
380    graph.load(graph_path, edge_loader);
381    graph.finalize();
382    
383    ///////////////////////////////////////////////////////
384    // Computer features in parallel
385    graph.transform_vertices(compute_features);
386
387    ///////////////////////////////////////////////////////
388    // Match features in parallel
389    graph.transform_edges(match_features);
390
391    
392    ///////////////////////////////////////////////////////
393    // Graphlab Engine
394    engine_type engine(dc, graph, clopts);
395
396    
397    ///////////////////////////////////////////////////////
398    // Run everything
399    engine.signal_all();
400    graphlab::timer timer;
401    engine.start();  
402    const double runtime = timer.current_time();
403    dc.cout() 
404    << "----------------------------------------------------------" << std::endl
405    << "Final Runtime (seconds):   " << runtime 
406    << std::endl
407    << "Updates executed: " << engine.num_updates() << std::endl
408    << "Update Rate (updates/second): " 
409    << engine.num_updates() / runtime << std::endl;
410    
411//    int retval = parseCmdArgs(argc, argv);
412//    if (retval)
413//        return retval;
414//    
415//    // Check if have enough images
416//    int num_images = static_cast<int>(img_names.size());
417//    if (num_images < 2)
418//    {
419//        LOGLN("Need more images");
420//        return -1;
421//    }
422//    
423//    double work_scale = 1, seam_scale = 1, compose_scale = 1;
424//    bool is_work_scale_set = false, is_seam_scale_set = false, is_compose_scale_set = false;
425    
426//    LOGLN("Finding features...");
427    
428}
429