From 590a5019f7fc28a17b8ba6203169d9de4a249325 Mon Sep 17 00:00:00 2001
From: Andrew Bell <andrew.bell.ia@gmail.com>
Date: Tue, 13 Mar 2018 15:15:17 -0400
Subject: [PATCH] Use raster types as PDAL types when creating dimensions.

---
 filters/ColorizationFilter.cpp | 12 +++++++++---
 filters/ColorizationFilter.hpp |  3 ++-
 io/GDALReader.cpp              |  7 +++----
 io/GDALReader.hpp              |  1 +
 4 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/filters/ColorizationFilter.cpp b/filters/ColorizationFilter.cpp
index 0aa804005e..2d2ee56b3a 100644
--- a/filters/ColorizationFilter.cpp
+++ b/filters/ColorizationFilter.cpp
@@ -131,6 +131,12 @@ void ColorizationFilter::addArgs(ProgramArgs& args)
 
 void ColorizationFilter::initialize()
 {
+    gdal::registerDrivers();
+
+    m_raster.reset(new gdal::Raster(m_rasterFilename));
+    auto bandTypes = m_raster->getPDALDimensionTypes();
+    m_raster->close();
+
     if (m_dimSpec.empty())
         m_dimSpec = { "Red", "Green", "Blue" };
 
@@ -141,6 +147,8 @@ void ColorizationFilter::initialize()
         {
             BandInfo bi = parseDim(dim, defaultBand);
             defaultBand = bi.m_band + 1;
+            if (bi.m_band < bandTypes.size())
+                bi.m_type = bandTypes[bi.m_band];
             m_bands.push_back(bi);
         }
         catch(const std::string& what)
@@ -149,15 +157,13 @@ void ColorizationFilter::initialize()
         }
     }
 
-    gdal::registerDrivers();
 }
 
 
 void ColorizationFilter::addDimensions(PointLayoutPtr layout)
 {
     for (auto& band : m_bands)
-        band.m_dim = layout->registerOrAssignDim(band.m_name,
-            Dimension::defaultType(Dimension::Id::Red));
+        band.m_dim = layout->registerOrAssignDim(band.m_name, band.m_type);
 }
 
 
diff --git a/filters/ColorizationFilter.hpp b/filters/ColorizationFilter.hpp
index 0e44e4b5c8..79ea71e69b 100644
--- a/filters/ColorizationFilter.hpp
+++ b/filters/ColorizationFilter.hpp
@@ -59,7 +59,7 @@ class PDAL_DLL ColorizationFilter : public Filter, public Streamable
     {
         BandInfo(const std::string& name, uint32_t band, double scale) :
             m_name(name), m_band(band), m_scale(scale),
-            m_dim(Dimension::Id::Unknown)
+            m_dim(Dimension::Id::Unknown), m_type(Dimension::Type::Double)
         {}
 
         BandInfo() : m_band(0), m_scale(1.0), m_dim(Dimension::Id::Unknown)
@@ -69,6 +69,7 @@ class PDAL_DLL ColorizationFilter : public Filter, public Streamable
         uint32_t m_band;
         double m_scale;
         Dimension::Id m_dim;
+        Dimension::Type m_type;
     };
 
 
diff --git a/io/GDALReader.cpp b/io/GDALReader.cpp
index 197edf510c..9f2366eb92 100644
--- a/io/GDALReader.cpp
+++ b/io/GDALReader.cpp
@@ -78,6 +78,7 @@ void GDALReader::initialize()
     }
 
     m_count = m_raster->width() * m_raster->height();
+    m_bandTypes = m_raster->getPDALDimensionTypes();
     m_raster->close();
 }
 
@@ -111,7 +112,7 @@ void GDALReader::addDimensions(PointLayoutPtr layout)
     {
         std::ostringstream oss;
         oss << "band-" << (i + 1);
-        layout->registerOrAssignDim(oss.str(), Dimension::Type::Double);
+        layout->registerOrAssignDim(oss.str(), m_bandTypes[i]);
     }
 }
 
@@ -142,13 +143,11 @@ point_count_t GDALReader::read(PointViewPtr view, point_count_t num)
     }
 
     std::vector<uint8_t> band;
-    std::vector<Dimension::Type> band_types =
-        m_raster->getPDALDimensionTypes();
 
     for (int b = 0; b < m_raster->bandCount(); ++b)
     {
         // Bands count from 1
-        switch (band_types[b])
+        switch (m_bandTypes[b])
         {
         case Dimension::Type::Signed8:
             readBandData<int8_t>(b + 1, view, count);
diff --git a/io/GDALReader.hpp b/io/GDALReader.hpp
index 9582b414f4..b84975ece1 100644
--- a/io/GDALReader.hpp
+++ b/io/GDALReader.hpp
@@ -72,6 +72,7 @@ class PDAL_DLL GDALReader : public Reader
     void readBandData(int band, PointViewPtr view, point_count_t count);
 
     std::unique_ptr<gdal::Raster> m_raster;
+    std::vector<Dimension::Type> m_bandTypes;
     point_count_t m_index;
 
 };