Add initial support to match histogram (brazil-data-cube#138)

cube_builder/celery/tasks.py

@@ -21,7 +21,7 @@
 from ..celery import celery_app
 from ..models import Activity
 from ..constants import CLEAR_OBSERVATION_NAME, TOTAL_OBSERVATION_NAME, PROVENANCE_NAME, DATASOURCE_NAME
-from ..utils.image import create_empty_raster
+from ..utils.image import create_empty_raster, match_histogram_with_merges
 from ..utils.processing import DataCubeFragments, build_cube_path, post_processing_quality
 from ..utils.processing import compute_data_set_stats, get_or_create_model
 from ..utils.processing import blend as blend_processing, merge as merge_processing, publish_datacube, publish_merge
@@ -277,6 +277,25 @@ def _is_not_stk(_merge):
 
         activities[_merge['band']] = activity
 
+    # TODO: Add option to skip histogram.
+    if kwargs.get('histogram_matching'):
+        ordered_best_efficacy = sorted(quality_date_stats.items(), key=lambda item: item[1][0], reverse=True)
+
+        best_date, (_, _, best_mask_file, _) = ordered_best_efficacy[0]
+        dates = map(lambda entry: entry[0], ordered_best_efficacy[1:])
+
+        for date in dates:
+            logging.info(f'Applying Histogram Matching: Reference date {best_date}, current {date}...')
+            for band, activity in activities.items():
+                reference = activities[band]['scenes'][best_date]['ARDfiles'][band]
+
+                if band == band_map['quality']:
+                    continue
+
+                source = activity['scenes'][date]['ARDfiles'][band]
+                source_mask = activity['scenes'][date]['ARDfiles'][band_map['quality']]
+                match_histogram_with_merges(source, source_mask, reference, best_mask_file)
+
     # Prepare list of activities to dispatch
     activity_list = list(activities.values())
 

cube_builder/cli.py

@@ -99,6 +99,7 @@ def worker(ctx: click.Context):
 @click.option('--with-rgb', is_flag=True, help='Generate a file with RGB bands, based in quick look.')
 @click.option('--token', type=click.STRING, help='Token to access data from STAC.')
 @click.option('--shape', type=click.STRING, help='Use custom output shape. i.e `--shape=10980x10980`')
+@click.option('--histogram-matching', is_flag=True, help='Match the histogram in the temporal composition function.')
 @with_appcontext
 def build(datacube: str, collections: str, tiles: str, start: str, end: str, bands: str = None,
           stac_url: str = None, force=False, with_rgb=False, shape=None, **kwargs):

cube_builder/forms.py

@@ -129,6 +129,7 @@ class DataCubeMetadataForm(Schema):
     title = fields.String(required=False, allow_none=False)
     public = fields.Boolean(required=False, allow_none=False, default=True)
 
+
 class DataCubeProcessForm(Schema):
     """Define parser for datacube generation."""
 
@@ -145,6 +146,7 @@ class DataCubeProcessForm(Schema):
     shape = fields.List(fields.Integer(required=False))
     # Reuse data cube from another data cube
     reuse_from = fields.String(required=False, allow_none=True)
+    histogram_matching = fields.Boolean(required=False, default=False)
 
 
 class PeriodForm(Schema):

cube_builder/utils/image.py

@@ -13,10 +13,11 @@
 from pathlib import Path
 from typing import List
 
+import numpy
 import rasterio
 from rasterio._warp import Affine
 from sqlalchemy.engine.result import ResultProxy, RowProxy
-from .processing import SmartDataSet, generate_cogs
+from .processing import SmartDataSet, generate_cogs, save_as_cog
 
 from ..config import Config
 
@@ -151,3 +152,53 @@ def create_empty_raster(location: str, proj4: str, dtype: str, xmin: float, ymax
         generate_cogs(str(location), str(location))
 
     return str(location)
+
+
+def match_histogram_with_merges(source: str, source_mask: str, reference: str, reference_mask: str):
+    """Normalize the source image histogram with reference image.
+
+    This functions implements the `skimage.exposure.match_histograms`, which consists in the manipulate the pixels of an
+    input image and match the histogram with the reference image.
+
+    See more in `Histogram Matching <https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_histogram_matching.html>`_.
+
+    Notes:
+        It overwrites the source file.
+
+    Args:
+        source (str): Path to the rasterio data set file
+        source_mask (str): Path to the rasterio data set file
+        reference (str): Path to the rasterio data set file
+        reference_mask (str): Path to the rasterio data set file
+    """
+    from skimage.exposure import match_histograms as _match_histograms
+
+    with rasterio.open(source) as source_data_set, rasterio.open(source_mask) as source_mask_data_set:
+        source_arr = source_data_set.read(1, masked=True)
+        source_mask_arr = source_mask_data_set.read(1)
+        source_options = source_data_set.profile.copy()
+
+    with rasterio.open(reference) as reference_data_set, rasterio.open(reference_mask) as reference_mask_data_set:
+        reference_arr = reference_data_set.read(1, masked=True)
+        reference_mask_arr = reference_mask_data_set.read(1)
+
+    intersect_mask = numpy.logical_and(
+        source_mask_arr < 255,  # CHECK: Use only valid data? numpy.isin(source_mask_arr, [0, 1, 3]),
+        reference_mask_arr < 255,   # CHECK: Use only valid data? numpy.isin(reference_mask_arr, [0, 1, 3]),
+    )
+
+    valid_positions = numpy.where(intersect_mask)
+
+    if valid_positions and len(valid_positions[0]) == 0:
+        return
+
+    intersected_source_arr = source_arr[valid_positions]
+    intersected_reference_arr = reference_arr[valid_positions]
+
+    histogram = _match_histograms(intersected_source_arr, intersected_reference_arr)
+
+    histogram = numpy.round(histogram).astype(source_options['dtype'])
+
+    source_arr[valid_positions] = histogram
+
+    save_as_cog(str(source), source_arr, mode='w', **source_options)

setup.py

@@ -32,9 +32,14 @@
     'check-manifest>=0.40',
 ]
 
+histogram_require = [
+    'scikit-image>=0.18,<1'
+]
+
 extras_require = {
     'docs': docs_require,
     'tests': tests_require,
+    'histogram': histogram_require
 }
 
 extras_require['all'] = [ req for exts, reqs in extras_require.items() for req in reqs ]