From 1bd0069948e8fb5dfb40a12a71607ca0a039095f Mon Sep 17 00:00:00 2001
From: Even Rouault <even.rouault@spatialys.com>
Date: Mon, 27 Nov 2023 14:52:01 +0100
Subject: [PATCH] COG: for JPEG compression, convert single band+alpha as
 single band JPEG + 1-bit mask band

Slightly related to https://github.com/OSGeo/gdal/issues/8834
---
 autotest/gcore/cog.py             | 21 +++++++++++++++++++++
 doc/source/drivers/raster/cog.rst |  2 ++
 frmts/gtiff/cogdriver.cpp         | 27 +++++++++++++++++++--------
 3 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/autotest/gcore/cog.py b/autotest/gcore/cog.py
index c07dd3e4744c..2b1cb961e686 100755
--- a/autotest/gcore/cog.py
+++ b/autotest/gcore/cog.py
@@ -295,6 +295,27 @@ def my_cbk(pct, _, arg):
     gdal.Unlink(directory)
 
 
+###############################################################################
+# Test creation from a single band + alpha dataset
+
+
+@pytest.mark.require_creation_option("COG", "JPEG")
+def test_cog_single_band_plus_alpha_jpeg_compression(tmp_vsimem):
+
+    filename = str(tmp_vsimem / "cog.tif")
+    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 2)
+    src_ds.GetRasterBand(2).SetColorInterpretation(gdal.GCI_AlphaBand)
+
+    ds = gdal.GetDriverByName("COG").CreateCopy(
+        filename,
+        src_ds,
+        options=["COMPRESS=JPEG"],
+    )
+
+    assert ds.RasterCount == 1
+    assert ds.GetRasterBand(1).GetMaskFlags() == gdal.GMF_PER_DATASET
+
+
 ###############################################################################
 # Test creation of overviews with a different compression method
 
diff --git a/doc/source/drivers/raster/cog.rst b/doc/source/drivers/raster/cog.rst
index a57de83a7b61..09bb961de23a 100644
--- a/doc/source/drivers/raster/cog.rst
+++ b/doc/source/drivers/raster/cog.rst
@@ -57,6 +57,8 @@ General creation options
         For the COG driver, JPEG compression for 3 or 4-band images automatically
         selects the PHOTOMETRIC=YCBCR colorspace with a 4:2:2 subsampling of the Y,Cb,Cr
         components.
+        For a input dataset (single-band or 3-band), plus an alpha band,
+        the alpha band will be converted as a 1-bit DEFLATE compressed mask.
 
       * ``LZW``, ``DEFLATE`` and ``ZSTD`` compressions can be used with the PREDICTOR creation option.
 
diff --git a/frmts/gtiff/cogdriver.cpp b/frmts/gtiff/cogdriver.cpp
index b0ef9cbdb943..77117799c52e 100644
--- a/frmts/gtiff/cogdriver.cpp
+++ b/frmts/gtiff/cogdriver.cpp
@@ -889,20 +889,31 @@ GDALDataset *GDALCOGCreator::Create(const char *pszFilename,
 
     CPLString osCompress = CSLFetchNameValueDef(papszOptions, "COMPRESS",
                                                 gbHasLZW ? "LZW" : "NONE");
-    if (EQUAL(osCompress, "JPEG") && poCurDS->GetRasterCount() == 4 &&
-        poCurDS->GetRasterBand(4)->GetColorInterpretation() == GCI_AlphaBand)
+    if (EQUAL(osCompress, "JPEG") &&
+        (poCurDS->GetRasterCount() == 2 || poCurDS->GetRasterCount() == 4) &&
+        poCurDS->GetRasterBand(poCurDS->GetRasterCount())
+                ->GetColorInterpretation() == GCI_AlphaBand)
     {
         char **papszArg = nullptr;
         papszArg = CSLAddString(papszArg, "-of");
         papszArg = CSLAddString(papszArg, "VRT");
         papszArg = CSLAddString(papszArg, "-b");
         papszArg = CSLAddString(papszArg, "1");
-        papszArg = CSLAddString(papszArg, "-b");
-        papszArg = CSLAddString(papszArg, "2");
-        papszArg = CSLAddString(papszArg, "-b");
-        papszArg = CSLAddString(papszArg, "3");
-        papszArg = CSLAddString(papszArg, "-mask");
-        papszArg = CSLAddString(papszArg, "4");
+        if (poCurDS->GetRasterCount() == 2)
+        {
+            papszArg = CSLAddString(papszArg, "-mask");
+            papszArg = CSLAddString(papszArg, "2");
+        }
+        else
+        {
+            CPLAssert(poCurDS->GetRasterCount() == 4);
+            papszArg = CSLAddString(papszArg, "-b");
+            papszArg = CSLAddString(papszArg, "2");
+            papszArg = CSLAddString(papszArg, "-b");
+            papszArg = CSLAddString(papszArg, "3");
+            papszArg = CSLAddString(papszArg, "-mask");
+            papszArg = CSLAddString(papszArg, "4");
+        }
         GDALTranslateOptions *psOptions =
             GDALTranslateOptionsNew(papszArg, nullptr);
         CSLDestroy(papszArg);