Adding latest changes to inference pipeline

README.md

@@ -97,6 +97,16 @@ Latest test:
 singularity exec --nv --env PYTHONPATH="$NOBACKUP/development/dl-water-bodies" -B $NOBACKUP,/explore/nobackup/people,/explore/nobackup/projects /explore/nobackup/projects/ilab/containers/vhr-cloudmask.sif python /explore/nobackup/people/jacaraba/development/dl-water-bodies/dl_water_bodies/view/dlwater_pipeline_cli.py -o '/explore/nobackup/projects/ilab/test/dlwater-test' -r '/explore/nobackup/people/almullen/smallsat_augmentation/data/planet/YKD/Ch009v024/Ch009v024_20210715_composite.tif' -s predict
 ```
 
+## Benchmark between versions
+
+This file has the subset of images we are testing with: /explore/nobackup/people/cssprad1/projects/planet_water/code/feature-composite-refactor/test_subset_202207.csv. The regex to cover these files is: /explore/nobackup/people/almullen/smallsat_augmentation/data/planet/YKD/Ch009v024/Ch009v024_202207*_composite.tif.
+
+For inference of the tiles:
+
+```bash
+singularity exec --nv --env PYTHONPATH="$NOBACKUP/development/dl-water-bodies" -B $NOBACKUP,/explore/nobackup/people,/explore/nobackup/projects /explore/nobackup/projects/ilab/containers/vhr-cloudmask.sif python /explore/nobackup/people/jacaraba/development/dl-water-bodies/dl_water_bodies/view/dlwater_pipeline_cli.py -o '/explore/nobackup/projects/ilab/test/dlwater-test' -r '/explore/nobackup/people/almullen/smallsat_augmentation/data/planet/YKD/Ch009v024/Ch009v024_202207*_composite.tif' -m '/explore/nobackup/people/almullen/smallsat_augmentation/data/planet_masks/YKD/Ch009v024/Ch009v024_202207*_composite_udm2.tif' -s predict
+```
+
 ## Predictor.py
 
 This python code handles lake and pond detection in PlanetScope imagery. To run the predictor you will first need to download the .h5 model files from https://doi.org/10.5281/zenodo.7682754. The script is setup to run on a directory of PlanetScope images and output predictions to another directory.

dl_water_bodies/config/dlwater_config.py

@@ -149,10 +149,17 @@ class WaterMaskConfig:
     # String regex to find rasters to predict
     inference_regex: Optional[str] = "*.tif"
 
+    # String regex to find rasters to mask predictions
+    mask_regex: Optional[str] = "*.tif"
+
     # List regex to find rasters to predict (multiple locations)
     inference_regex_list: Optional[List[str]] = field(
         default_factory=lambda: [])
 
+    # List regex to find rasters to mask predictions (multiple locations)
+    mask_regex_list: Optional[List[str]] = field(
+        default_factory=lambda: [])
+
     # Window size for sliding window operations
     window_size: Optional[int] = 8120
 

dl_water_bodies/pipelines/dlwater_pipeline.py

@@ -16,7 +16,9 @@
 from multiprocessing import Pool, cpu_count
 from omegaconf.listconfig import ListConfig
 
-from tensorflow_caney.model.config.cnn_config import Config
+# from tensorflow_caney.model.config.cnn_config import Config
+from dl_water_bodies.config.dlwater_config \
+    import WaterMaskConfig as Config
 from tensorflow_caney.utils.system import seed_everything, \
     set_gpu_strategy, set_mixed_precision, set_xla, array_module
 from tensorflow_caney.utils.data import read_dataset_csv, \
@@ -61,6 +63,7 @@ def __init__(
                 model_filename: str = None,
                 output_dir: str = None,
                 inference_regex_list: str = None,
+                mask_regex_list: str = None,
                 force_delete: bool = False,
                 default_config: str = 'templates/watermask_default.yaml',
                 logger=None
@@ -97,6 +100,10 @@ def __init__(
         if inference_regex_list is not None:
             self.conf.inference_regex_list = inference_regex_list
 
+        # rewrite mask regex list
+        if mask_regex_list is not None:
+            self.conf.mask_regex_list = mask_regex_list
+
         # Set Data CSV
         self.data_csv = data_csv
 
@@ -212,13 +219,23 @@ def predict(self) -> None:
 
         # Gather filenames to predict
         if len(self.conf.inference_regex_list) > 0:
-            data_filenames = self.get_filenames(self.conf.inference_regex_list)
+            data_filenames = sorted(
+                self.get_filenames(self.conf.inference_regex_list))
         else:
-            data_filenames = self.get_filenames(self.conf.inference_regex)
+            data_filenames = sorted(
+                self.get_filenames(self.conf.inference_regex))
         logging.info(f'{len(data_filenames)} files to predict')
 
+        # Gather filenames to mask predictions
+        if len(self.conf.mask_regex_list) > 0:
+            mask_filenames = sorted(
+                self.get_filenames(self.conf.mask_regex_list))
+        else:
+            mask_filenames = sorted(self.get_filenames(self.conf.mask_regex))
+        logging.info(f'{len(mask_filenames)} masks for prediction')
+
         # iterate files, create lock file to avoid predicting the same file
-        for filename in sorted(data_filenames):
+        for filename, mask_filename in zip(data_filenames, mask_filenames):
 
             # start timer
             start_time = time.time()
@@ -252,10 +269,14 @@ def predict(self) -> None:
                 # create lock file
                 open(lock_filename, 'w').close()
 
-                # open filename
+                # open data filename
                 image = rxr.open_rasterio(filename)
                 logging.info(f'Prediction shape: {image.shape}')
 
+                # open mask filename
+                mask = rxr.open_rasterio(mask_filename)
+                logging.info(f'Mask shape: {mask.shape}')
+
                 # check bands in imagery, do not proceed if one band
                 if image.shape[0] == 1:
                     logging.info(
@@ -273,7 +294,6 @@ def predict(self) -> None:
 
             # scale image values, PlanetScope specific normalization
             image = (image / 10000.0) * 255
-
             temporary_tif = preprocess_input(image.values)
 
             # Sliding window prediction
@@ -294,6 +314,19 @@ def predict(self) -> None:
                     probability_map=self.conf.probability_map
                 )
 
+            # Set nodata values on mask
+            # Planet imagery does not have the burned no-data.
+            # The default is 0 which we are adding here.
+            nodata = 0  # prediction.rio.nodata
+
+            # Mask out using the Planet quota
+            mask = np.squeeze(mask[0, :, :].values)
+            prediction[mask == 0] = nodata
+
+            # TODO: Fix no-data from corners to 255
+            # nodata_mask = np.squeeze(image[:, :, 0].values)
+            # prediction[nodata_mask == 0] = np.nan#255
+
             # Drop image band to allow for a merge of mask
             image = image.drop(
                 dim="band",
@@ -314,23 +347,12 @@ def predict(self) -> None:
             prediction.attrs['model_name'] = (self.conf.model_filename)
             prediction = prediction.transpose("band", "y", "x")
 
-            # Set nodata values on mask
-            # Planet imagery does not have the burned no-data.
-            # The default is 0 which we are adding here.
-            nodata = 0  # prediction.rio.nodata
-
-            print('The prediction no-data', nodata)
-            print('Before the where', np.unique(prediction.values))
-
+            # do not need this if it comes from numpy
             prediction = prediction.where(image != nodata)
-            print('After the where', np.unique(prediction.values))
-            print(f"nodata: {prediction.rio.nodata}")
-            print(f"encoded_nodata: {prediction.rio.encoded_nodata}")
-
-            prediction.rio.write_nodata(
-                self.conf.prediction_nodata, encoded=True, inplace=True)
-            print(f"nodata: {prediction.rio.nodata}")
-            print(f"encoded_nodata: {prediction.rio.encoded_nodata}")
+
+            # prediction.rio.write_nodata(
+            #    self.conf.prediction_nodata, encoded=True, inplace=True)
+            prediction.rio.write_nodata(nodata, encoded=True, inplace=True)
 
             # Save output raster file to disk
             prediction.rio.to_raster(

dl_water_bodies/view/dlwater_pipeline_cli.py

@@ -67,6 +67,15 @@ def main():
                         help='Inference regex list',
                         default=['*.tif'])
 
+    parser.add_argument('-mr',
+                        '--mask-regex-list',
+                        type=str,
+                        nargs='*',
+                        required=False,
+                        dest='mask_regex_list',
+                        help='Mask regex list',
+                        default=['*.tif'])
+
     parser.add_argument('-f',
                         '--force-delete',
                         default=False,
@@ -87,6 +96,7 @@ def main():
         args.model_filename,
         args.output_dir,
         args.inference_regex_list,
+        args.mask_regex_list,
         args.force_delete
     )