From aafdb32d71e7adc4182943cd9cda87155719e561 Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Wed, 22 Jul 2020 18:35:25 -0400
Subject: [PATCH] Started writing GDAL based COG translation

---
 app/cogeo.py | 69 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 67 insertions(+), 2 deletions(-)

diff --git a/app/cogeo.py b/app/cogeo.py
index 3ec30d01..9cb7ab1d 100644
--- a/app/cogeo.py
+++ b/app/cogeo.py
@@ -3,6 +3,8 @@ import logging
 import tempfile
 import shutil
 import rasterio
+import re
+import subprocess
 from rio_cogeo.cogeo import cog_validate, cog_translate
 from rio_tiler.utils import has_alpha_band
 from webodm import settings
@@ -18,12 +20,13 @@ def valid_cogeo(src_path):
     return cog_validate(src_path, strict=True)
 
 
-def assure_cogeo(src_path):
+def assure_cogeo(src_path, use_legacy=False):
     """
     Guarantee that the .tif passed as an argument is a Cloud Optimized GeoTIFF (cogeo)
     If the path is not a cogeo, it is destructively converted into a cogeo.
     If the file cannot be converted, the function does not change the file
     :param src_path: path to GeoTIFF (cogeo or not)
+    :param use_legacy: whether to force the use of legacy implementation. By default the best implementation is used.
     :return: None
     """
 
@@ -36,6 +39,64 @@ def assure_cogeo(src_path):
 
     # Not a cogeo
     logger.info("Optimizing %s as Cloud Optimized GeoTIFF" % src_path)
+
+    # Check if we have GDAL >= 3.1
+    use_legacy = False
+    gdal_version = get_gdal_version()
+    if gdal_version:
+        major, minor, build = gdal_version
+        
+        # GDAL 2 and lower
+        if major <= 2:
+            use_legacy = True
+        
+        # GDAL 3.0 and lower
+        if major == 3 and minor < 1:
+            use_legacy = True
+    else:
+        # This shouldn't happen
+        use_legacy = True
+        
+    if use_legacy:
+        logger.info("Using legacy implementation (GDAL >= 3.1 not found)")
+        make_cogeo_legacy(src_path)
+    else:
+        make_cogeo_gdal(src_path)
+
+def get_gdal_version():
+    # Bit of a hack without installing 
+    # python bindings
+    gdal_translate = shutil.which('gdal_translate')
+    if not gdal_translate:
+        return None
+    
+    # Get version
+    version_output = subprocess.check_output("%s --version" % gdal_translate)
+    
+    m = re.match(r"GDAL\s+([\d+])\.([\d+])\.([\d+]),\s+released", version_output)
+    if not m:
+        return None
+    
+    return tuple(map(int, m.groups()))
+
+
+def make_cogeo_gdal(src_path):
+    """
+    Make src_path a Cloud Optimized GeoTIFF.
+    Requires GDAL >= 3.1
+    """
+
+    tmpfile = tempfile.mktemp('_cogeo.tif', dir=settings.MEDIA_TMP)
+    swapfile = tempfile.mktemp('_cogeo_swap.tif', dir=settings.MEDIA_TMP)
+
+    # gdal_translate -of COG -co BLOCKSIZE=256 -co COMPRESS=deflate -co NUM_THREADS=4 -co BIGTIFF=IF_SAFER -co QUALITY=100 -co SPARSE_OK=ON --config GDAL_NUM_THREADS ALL_CPUS brighton.tif cog.tif
+
+def make_cogeo_legacy(src_path):
+    """
+    Make src_path a Cloud Optimized GeoTIFF
+    This implementation does not require GDAL >= 3.1
+    but sometimes (rarely) hangs for unknown reasons
+    """
     tmpfile = tempfile.mktemp('_cogeo.tif', dir=settings.MEDIA_TMP)
     swapfile = tempfile.mktemp('_cogeo_swap.tif', dir=settings.MEDIA_TMP)
 
@@ -77,4 +138,8 @@ def assure_cogeo(src_path):
             raise e
 
         if os.path.isfile(swapfile):
-            os.remove(swapfile)
\ No newline at end of file
+            os.remove(swapfile)
+
+        return True
+    else:
+        return False
\ No newline at end of file