sandboxed-api/oss-internship-2020/gdal/raster_to_gtiff/get_raster_data.cc

// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "get_raster_data.h"  // NOLINT(build/include)

#include <memory>

#include "gdal.h"

namespace gdal::sandbox::parser {

namespace {

inline constexpr int kGeoTransformSize = 6;

}  // namespace

RasterDataset GetRasterBandsFromFile(const std::string& filename) {
  GDALAllRegister();
  GDALDatasetH dataset = GDALOpen(filename.data(), GA_ReadOnly);
  GDALDriverH driver = GDALGetDatasetDriver(dataset);

  RasterDataset result_dataset = {
    GDALGetRasterXSize(dataset),
    GDALGetRasterYSize(dataset)
  };

  if (GDALGetProjectionRef(dataset) != nullptr) {
    result_dataset.wkt_projection = std::string(GDALGetProjectionRef(dataset));
  }

  std::vector<double> geo_transform(kGeoTransformSize, 0.0);

  if (GDALGetGeoTransform(dataset, geo_transform.data()) == CE_None) {
    result_dataset.geo_transform = std::move(geo_transform);
  }

  int bands_count = GDALGetRasterCount(dataset);

  std::vector<RasterBandData> bands_data;
  bands_data.reserve(bands_count);

  for (int i = 1; i <= bands_count; ++i) {
    GDALRasterBandH band = GDALGetRasterBand(dataset, i);
    int width = GDALGetRasterBandXSize(band);
    int height = GDALGetRasterBandYSize(band);

    std::unique_ptr<int> no_data_result = nullptr;
    double no_data_value = GDALGetRasterNoDataValue(band, no_data_result.get());
    std::optional<double> no_data_value_holder = 
        no_data_result.get() == nullptr ? std::nullopt 
          : std::make_optional<double>(no_data_value);

    int data_type  = static_cast<int>(GDALGetRasterDataType(band));
    int color_interp = static_cast<int>(GDALGetRasterColorInterpretation(band));

    std::vector<int> band_raster_data;
    band_raster_data.resize(width * height);

    // GDALRasterIO with GF_Write should use the same type (GDT_Int32)
    GDALRasterIO(band, GF_Read, 0, 0, width, height, band_raster_data.data(),
        width, height, GDT_Int32, 0, 0);

    bands_data.push_back(
      {
        width,
        height,
        std::move(band_raster_data),
        data_type,
        color_interp,
        std::move(no_data_value_holder)
      }
    );
  }
  
  result_dataset.bands = std::move(bands_data);

  GDALClose(dataset);

  return result_dataset;
}

bool operator==(const RasterBandData& lhs, const RasterBandData& rhs) {
  return lhs.width == rhs.width && lhs.height == rhs.height
          && lhs.data == rhs.data && lhs.data_type == rhs.data_type 
          && lhs.color_interp == rhs.color_interp
          && lhs.no_data_value == rhs.no_data_value;
}

bool operator==(const RasterDataset& lhs, const RasterDataset& rhs) {
  return lhs.width == rhs.width && lhs.height == rhs.height
          && lhs.bands == rhs.bands 
          && lhs.wkt_projection == rhs.wkt_projection
          && lhs.geo_transform == rhs.geo_transform;
}

}  // namespace gdal::sandbox::parser