From c7b171d690fa85af2ab0f7fbd81208ff997215a4 Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Mon, 23 Dec 2019 09:16:54 -0500
Subject: [PATCH] Option to set RGB band color interpolation in ortho

Former-commit-id: 826b04b7c4590829a7c613d5dbe2237fe92732be
---
 modules/odm_orthophoto/src/OdmOrthoPhoto.cpp | 39 +++++++++++++++-----
 modules/odm_orthophoto/src/OdmOrthoPhoto.hpp |  2 +
 stages/odm_orthophoto.py                     |  4 +-
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/modules/odm_orthophoto/src/OdmOrthoPhoto.cpp b/modules/odm_orthophoto/src/OdmOrthoPhoto.cpp
index 370dc7c7..d4dd905f 100644
--- a/modules/odm_orthophoto/src/OdmOrthoPhoto.cpp
+++ b/modules/odm_orthophoto/src/OdmOrthoPhoto.cpp
@@ -10,6 +10,7 @@ OdmOrthoPhoto::OdmOrthoPhoto()
     outputFile_ = "ortho.tif";
     logFile_    = "log.txt";
     outputCornerFile_ = "";
+    bandsOrder = "red,green,blue";
 
     resolution_ = 0.0f;
 
@@ -127,6 +128,15 @@ void OdmOrthoPhoto::parseArguments(int argc, char *argv[])
                 inputFiles.push_back(item);
             }
         }
+        else if(argument == "-bands")
+        {
+            argIndex++;
+            if (argIndex >= argc)
+            {
+                throw OdmOrthoPhotoException("Argument '" + argument + "' expects 1 more input following it, but no more inputs were provided.");
+            }
+            bandsOrder = std::string(argv[argIndex]);
+        }
         else if(argument == "-outputFile")
         {
             argIndex++;
@@ -154,6 +164,20 @@ void OdmOrthoPhoto::parseArguments(int argc, char *argv[])
         }
     }
     log_ << "\n";
+
+    std::stringstream ss(bandsOrder);
+    std::string item;
+    while(std::getline(ss, item, ',')){
+        if (item == "red" || item == "r"){
+            colorInterps.push_back(GCI_RedBand);
+        }else if (item == "green" || item == "g"){
+            colorInterps.push_back(GCI_GreenBand);
+        }else if (item == "blue" || item == "b"){
+            colorInterps.push_back(GCI_BlueBand);
+        }else{
+            colorInterps.push_back(GCI_GrayIndex);
+        }
+    }
 }
 
 void OdmOrthoPhoto::printHelp()
@@ -185,6 +209,9 @@ void OdmOrthoPhoto::printHelp()
     log_ << "\"-resolution <pixels/m>\" (mandatory)\n";
     log_ << "\"The number of pixels used per meter.\n\n";
 
+    log_ << "\"-bands red,green,blue,[...]\" (optional)\n";
+    log_ << "\"Naming of bands to assign color interpolation values when creating output TIFF.\n\n";
+
     log_.setIsPrintingInCout(false);
 }
 
@@ -206,15 +233,9 @@ void OdmOrthoPhoto::saveTIFF(const std::string &filename, GDALDataType dataType)
         hBand = GDALGetRasterBand( hDstDS, static_cast<int>(i) + 1 );
 
         // TODO: should we set these based on a command line parameter?
-        GDALColorInterp interp;
-        if (i == 0){
-            interp = GCI_RedBand;
-        }else if (i == 1){
-            interp = GCI_GreenBand;
-        }else if (i == 2){
-            interp = GCI_BlueBand;
-        }else{
-            interp = GCI_GrayIndex;
+        GDALColorInterp interp = GCI_GrayIndex;
+        if (i < colorInterps.size()){
+            interp = colorInterps[i];
         }
         GDALSetRasterColorInterpretation(hBand, interp );
 
diff --git a/modules/odm_orthophoto/src/OdmOrthoPhoto.hpp b/modules/odm_orthophoto/src/OdmOrthoPhoto.hpp
index b902b32f..e39c3529 100644
--- a/modules/odm_orthophoto/src/OdmOrthoPhoto.hpp
+++ b/modules/odm_orthophoto/src/OdmOrthoPhoto.hpp
@@ -177,10 +177,12 @@ private:
     std::string     outputFile_;        /**< Path to the destination file. */
     std::string     outputCornerFile_;  /**< Path to the output corner file. */
     std::string     logFile_;           /**< Path to the log file. */
+    std::string     bandsOrder;
 
     float           resolution_;        /**< The number of pixels per meter in the ortho photo. */
 
     std::vector<void *>    bands;
+    std::vector<GDALColorInterp> colorInterps;
     void *alphaBand; // Keep alpha band separate
     int currentBandIndex;
 
diff --git a/stages/odm_orthophoto.py b/stages/odm_orthophoto.py
index 623964bf..d5b8ed34 100644
--- a/stages/odm_orthophoto.py
+++ b/stages/odm_orthophoto.py
@@ -29,6 +29,7 @@ class ODMOrthoPhotoStage(types.ODM_Stage):
                 'ortho': tree.odm_orthophoto_render,
                 'corners': tree.odm_orthophoto_corners,
                 'res': 1.0 / (gsd.cap_resolution(args.orthophoto_resolution, tree.opensfm_reconstruction, ignore_gsd=args.ignore_gsd) / 100.0),
+                'bands': '',
                 'verbose': verbose
             }
 
@@ -63,6 +64,7 @@ class ODMOrthoPhotoStage(types.ODM_Stage):
                     if not primary:
                         subdir = band['name'].lower()
                     models.append(os.path.join(base_dir, subdir, model_file))
+                kwargs['bands'] = '-bands %s' % (','.join([quote(b['name'].lower()) for b in reconstruction.multi_camera]))
             else:
                 models.append(os.path.join(base_dir, model_file))
 
@@ -71,7 +73,7 @@ class ODMOrthoPhotoStage(types.ODM_Stage):
             # run odm_orthophoto
             system.run('{bin}/odm_orthophoto -inputFiles {models} '
                        '-logFile {log} -outputFile {ortho} -resolution {res} {verbose} '
-                       '-outputCornerFile {corners}'.format(**kwargs))
+                       '-outputCornerFile {corners} {bands}'.format(**kwargs))
 
             # Create georeferenced GeoTiff
             geotiffcreated = False