Made timestamps flexible and updated description

app.py

@@ -1,6 +1,14 @@
-import logging
+
 import os
+import torch
+import yaml
+import numpy as np
+import gradio as gr
+from einops import rearrange
+from functools import partial
 from huggingface_hub import hf_hub_download
+from prithvi_mae import PrithviMAE
+from inference import process_channel_group, _convert_np_uint8, load_example, run_model
 
 # pull files from hub
 token = os.environ.get("HF_TOKEN", None)
@@ -15,206 +23,6 @@ model_inference = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-2.0-30
 os.system(f'cp {model_def} .')
 os.system(f'cp {model_inference} .')
 
-import os
-import torch
-import yaml
-import numpy as np
-import gradio as gr
-from einops import rearrange
-from functools import partial
-from prithvi_mae import PrithviMAE
-from inference import process_channel_group, read_geotiff, save_geotiff, _convert_np_uint8, load_example, run_model
-
-
-NO_DATA = -9999
-NO_DATA_FLOAT = 0.0001
-PERCENTILES = (0.1, 99.9)
-
-
-# def process_channel_group(orig_img, new_img, channels, data_mean, data_std):
-#     """ Process *orig_img* and *new_img* for RGB visualization. Each band is rescaled back to the
-#         original range using *data_mean* and *data_std* and then lowest and highest percentiles are
-#         removed to enhance contrast. Data is rescaled to (0, 1) range and stacked channels_first.
-#     Args:
-#         orig_img: torch.Tensor representing original image (reference) with shape = (bands, H, W).
-#         new_img: torch.Tensor representing image with shape = (bands, H, W).
-#         channels: list of indices representing RGB channels.
-#         data_mean: list of mean values for each band.
-#         data_std: list of std values for each band.
-#     Returns:
-#         torch.Tensor with shape (num_channels, height, width) for original image
-#         torch.Tensor with shape (num_channels, height, width) for the other image
-#     """
-# 
-#     stack_c = [], []
-# 
-#     for c in channels:
-#         orig_ch = orig_img[c, ...]
-#         valid_mask = torch.ones_like(orig_ch, dtype=torch.bool)
-#         valid_mask[orig_ch == NO_DATA_FLOAT] = False
-# 
-#         # Back to original data range
-#         orig_ch = (orig_ch * data_std[c]) + data_mean[c]
-#         new_ch = (new_img[c, ...] * data_std[c]) + data_mean[c]
-# 
-#         # Rescale (enhancing contrast)
-#         min_value, max_value = np.percentile(orig_ch[valid_mask], PERCENTILES)
-# 
-#         orig_ch = torch.clamp((orig_ch - min_value) / (max_value - min_value), 0, 1)
-#         new_ch = torch.clamp((new_ch - min_value) / (max_value - min_value), 0, 1)
-# 
-#         # No data as zeros
-#         orig_ch[~valid_mask] = 0
-#         new_ch[~valid_mask] = 0
-# 
-#         stack_c[0].append(orig_ch)
-#         stack_c[1].append(new_ch)
-# 
-#     # Channels first
-#     stack_orig = torch.stack(stack_c[0], dim=0)
-#     stack_rec = torch.stack(stack_c[1], dim=0)
-# 
-#     return stack_orig, stack_rec
-# 
-# 
-# def read_geotiff(file_path: str):
-#     """ Read all bands from *file_path* and returns image + meta info.
-#     Args:
-#         file_path: path to image file.
-#     Returns:
-#         np.ndarray with shape (bands, height, width)
-#         meta info dict
-#     """
-# 
-#     with rasterio.open(file_path) as src:
-#         img = src.read()
-#         meta = src.meta
-#         coords = src.lnglat()
-# 
-#     return img, meta, coords
-# 
-# 
-# def save_geotiff(image, output_path: str, meta: dict):
-#     """ Save multi-band image in Geotiff file.
-#     Args:
-#         image: np.ndarray with shape (bands, height, width)
-#         output_path: path where to save the image
-#         meta: dict with meta info.
-#     """
-# 
-#     with rasterio.open(output_path, "w", **meta) as dest:
-#         for i in range(image.shape[0]):
-#             dest.write(image[i, :, :], i + 1)
-# 
-#     return
-# 
-# 
-# def _convert_np_uint8(float_image: torch.Tensor):
-# 
-#     image = float_image.numpy() * 255.0
-#     image = image.astype(dtype=np.uint8)
-#     image = image.transpose((1, 2, 0))
-# 
-#     return image
-# 
-# 
-# def load_example(file_paths: List[str], mean: List[float], std: List[float]):
-#     """ Build an input example by loading images in *file_paths*.
-#     Args:
-#         file_paths: list of file paths .
-#         mean: list containing mean values for each band in the images in *file_paths*.
-#         std: list containing std values for each band in the images in *file_paths*.
-#     Returns:
-#         np.array containing created example
-#         list of meta info for each image in *file_paths*
-#     """
-# 
-#     imgs = []
-#     metas = []
-# 
-#     for file in file_paths:
-#         img, meta = read_geotiff(file)
-#         img = img[:6]*10000 if img[:6].mean() <= 2 else img[:6]
-# 
-#         # Rescaling (don't normalize on nodata)
-#         img = np.moveaxis(img, 0, -1)   # channels last for rescaling
-#         img = np.where(img == NO_DATA, NO_DATA_FLOAT, (img - mean) / std)
-# 
-#         imgs.append(img)
-#         metas.append(meta)
-# 
-#     imgs = np.stack(imgs, axis=0)    # num_frames, H, W, C
-#     imgs = np.moveaxis(imgs, -1, 0).astype('float32')  # C, num_frames, H, W
-#     imgs = np.expand_dims(imgs, axis=0)  # add batch dim
-# 
-#     return imgs, metas
-# 
-# 
-# def run_model(model: torch.nn.Module, input_data: torch.Tensor, mask_ratio: float, device: torch.device):
-#     """ Run *model* with *input_data* and create images from output tokens (mask, reconstructed + visible).
-#     Args:
-#         model: MAE model to run.
-#         input_data: torch.Tensor with shape (B, C, T, H, W).
-#         mask_ratio: mask ratio to use.
-#         device: device where model should run.
-#     Returns:
-#         3 torch.Tensor with shape (B, C, T, H, W).
-#     """
-# 
-#     with torch.no_grad():
-#         x = input_data.to(device)
-# 
-#         _, pred, mask = model(x, mask_ratio)
-# 
-#     # Create mask and prediction images (un-patchify)
-#     mask_img = model.unpatchify(mask.unsqueeze(-1).repeat(1, 1, pred.shape[-1])).detach().cpu()
-#     pred_img = model.unpatchify(pred).detach().cpu()
-# 
-#     # Mix visible and predicted patches
-#     rec_img = input_data.clone()
-#     rec_img[mask_img == 1] = pred_img[mask_img == 1]  # binary mask: 0 is keep, 1 is remove
-# 
-#     # Switch zeros/ones in mask images so masked patches appear darker in plots (better visualization)
-#     mask_img = (~(mask_img.to(torch.bool))).to(torch.float)
-# 
-#     return rec_img, mask_img
-# 
-# 
-# def save_rgb_imgs(input_img, rec_img, mask_img, channels, mean, std, output_dir, meta_data):
-#     """ Wrapper function to save Geotiff images (original, reconstructed, masked) per timestamp.
-#     Args:
-#         input_img: input torch.Tensor with shape (C, T, H, W).
-#         rec_img: reconstructed torch.Tensor with shape (C, T, H, W).
-#         mask_img: mask torch.Tensor with shape (C, T, H, W).
-#         channels: list of indices representing RGB channels.
-#         mean: list of mean values for each band.
-#         std: list of std values for each band.
-#         output_dir: directory where to save outputs.
-#         meta_data: list of dicts with geotiff meta info.
-#     """
-# 
-#     for t in range(input_img.shape[1]):
-#         rgb_orig, rgb_pred = process_channel_group(orig_img=input_img[:, t, :, :],
-#                                                    new_img=rec_img[:, t, :, :],
-#                                                    channels=channels, data_mean=mean,
-#                                                    data_std=std)
-# 
-#         rgb_mask = mask_img[channels, t, :, :] * rgb_orig
-# 
-#         # Saving images
-# 
-#         save_geotiff(image=_convert_np_uint8(rgb_orig),
-#                      output_path=os.path.join(output_dir, f"original_rgb_t{t}.tiff"),
-#                      meta=meta_data[t])
-# 
-#         save_geotiff(image=_convert_np_uint8(rgb_pred),
-#                      output_path=os.path.join(output_dir, f"predicted_rgb_t{t}.tiff"),
-#                      meta=meta_data[t])
-# 
-#         save_geotiff(image=_convert_np_uint8(rgb_mask),
-#                      output_path=os.path.join(output_dir, f"masked_rgb_t{t}.tiff"),
-#                      meta=meta_data[t])
-
 
 def extract_rgb_imgs(input_img, rec_img, mask_img, channels, mean, std):
     """ Wrapper function to save Geotiff images (original, reconstructed, masked) per timestamp.
@@ -245,15 +53,21 @@ def extract_rgb_imgs(input_img, rec_img, mask_img, channels, mean, std):
         rgb_orig_list.append(_convert_np_uint8(rgb_orig).transpose(1, 2, 0))
         rgb_mask_list.append(_convert_np_uint8(rgb_mask).transpose(1, 2, 0))
         rgb_pred_list.append(_convert_np_uint8(rgb_pred).transpose(1, 2, 0))
-        
+
+    # Add white dummy image values for missing timestamps
+    dummy = np.ones((20, 20), dtype=np.uint8) * 255
+    num_dummies = 4 - len(rgb_orig_list)
+    if num_dummies:
+        rgb_orig_list.extend([dummy] * num_dummies)
+        rgb_mask_list.extend([dummy] * num_dummies)
+        rgb_pred_list.extend([dummy] * num_dummies)
+
     outputs = rgb_orig_list + rgb_mask_list + rgb_pred_list
 
     return outputs
 
 
-def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str, checkpoint: str):
-
-    
+def predict_on_images(data_files: list, yaml_file_path: str, checkpoint: str, mask_ratio: float = None):
     try:
         data_files = [x.name for x in data_files]
         print('Path extracted from example')
@@ -276,9 +90,7 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
 
     mask_ratio = mask_ratio or config['DATA']['MASK_RATIO']
 
-    if num_frames > 4:
-        # TODO: Check if we can limit this via UI
-        logging.warning("Model was only trained with only four timestamps.")
+    assert num_frames <= 4, "Demo only supports up to four timestamps"
 
     if torch.cuda.is_available():
         device = torch.device('cuda')
@@ -384,7 +196,7 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
     return outputs
 
 
-func = partial(predict_on_images, yaml_file_path=yaml_file_path,checkpoint=checkpoint)
+run_inference = partial(predict_on_images, yaml_file_path=yaml_file_path,checkpoint=checkpoint)
 
 with gr.Blocks() as demo:
 
@@ -397,10 +209,14 @@ Additionally, temporal and location information is added to the model input via
 More info about the model are available [here](https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL).\n
 
 This demo showcases the image reconstruction over one to four timestamps. 
-The model randomly masks out some proportion of the images and then reconstructing them based on the not masked portion of the images.\n
+The model randomly masks out some proportion of the images and reconstructs them based on the not masked portion of the images. 
+The reconstructed images are merged with the visible unmasked patches. 
+We recommend submitting images of size 224 to ~1000 pixels for faster processing time. 
+Images bigger than 224x224 are processed using a sliding window approach which can lead to artefacts between patches.\n
+
 The user needs to provide the HLS geotiff images, including the following channels in reflectance units: Blue, Green, Red, Narrow NIR, SWIR, SWIR 2. 
 Optionally, the location information is extracted from the tif files while the temporal information can be provided in the filename in the format `<date>T<time>` or `<year><julian day>T<time>` (HLS format). 
-We recommend submitting images of size 224 to 1000 pixels for faster processing time. Some example images are provided at the end of this page.
+Some example images are provided at the end of this page.
 ''')
     with gr.Row():
         with gr.Column():
@@ -409,48 +225,26 @@ We recommend submitting images of size 224 to 1000 pixels for faster processing
             btn = gr.Button("Submit")
     with gr.Row():
         gr.Markdown(value='## Original images')
-    with gr.Row():
-        gr.Markdown(value='T1')
-        gr.Markdown(value='T2')
-        gr.Markdown(value='T3')
-    with gr.Row():
-        out1_orig_t1 = gr.Image(image_mode='RGB')
-        out2_orig_t2 = gr.Image(image_mode='RGB')
-        out3_orig_t3 = gr.Image(image_mode='RGB')
-    with gr.Row():
         gr.Markdown(value='## Masked images')
-    with gr.Row():
-        gr.Markdown(value='T1')
-        gr.Markdown(value='T2')
-        gr.Markdown(value='T3')
-    with gr.Row():
-        out4_masked_t1 = gr.Image(image_mode='RGB')
-        out5_masked_t2 = gr.Image(image_mode='RGB')
-        out6_masked_t3 = gr.Image(image_mode='RGB')
-    with gr.Row():
-        gr.Markdown(value='## Reonstructed images')
-    with gr.Row():
-        gr.Markdown(value='T1')
-        gr.Markdown(value='T2')
-        gr.Markdown(value='T3')
-    with gr.Row():
-        out7_pred_t1 = gr.Image(image_mode='RGB')
-        out8_pred_t2 = gr.Image(image_mode='RGB')
-        out9_pred_t3 = gr.Image(image_mode='RGB')
-
-
-    btn.click(fn=func, 
-              # inputs=[inp_files, inp_slider], 
+        gr.Markdown(value='## Visible and reconstructed images')
+
+    original = []
+    masked = []
+    predicted = []
+    timestamps = []
+    for t in range(4):
+        timestamps.append(gr.Column(visible=t == 0))
+        with timestamps[t]:
+            with gr.Row():
+                gr.Markdown(value=f"Timestamp {t+1}")
+            with gr.Row():
+                original.append(gr.Image(image_mode='RGB'))
+                masked.append(gr.Image(image_mode='RGB'))
+                predicted.append(gr.Image(image_mode='RGB'))
+
+    btn.click(fn=run_inference,
               inputs=inp_files,
-              outputs=[out1_orig_t1, 
-                       out2_orig_t2, 
-                       out3_orig_t3, 
-                       out4_masked_t1, 
-                       out5_masked_t2, 
-                       out6_masked_t3,
-                       out7_pred_t1,
-                       out8_pred_t2,
-                       out9_pred_t3])
+              outputs=original + masked + predicted)
 
     with gr.Row():
         gr.Examples(examples=[[[os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T13REN.2018013T172747.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
@@ -463,17 +257,16 @@ We recommend submitting images of size 224 to 1000 pixels for faster processing
                       os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T18TVL.2018141T154435.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
                       os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T18TVL.2018189T154446.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif")]]],
                     inputs=inp_files,
-                    outputs=[out1_orig_t1,
-                           out2_orig_t2,
-                           out3_orig_t3,
-                           out4_masked_t1,
-                           out5_masked_t2,
-                           out6_masked_t3,
-                           out7_pred_t1,
-                           out8_pred_t2,
-                           out9_pred_t3],
-                    fn=func,
+                    outputs=original + masked + predicted,
+                    fn=run_inference,
                     cache_examples=True
     )
 
+    def update_visibility(files):
+        timestamps = [gr.Column(visible=t < len(files)) for t in range(4)]
+
+        return timestamps
+
+    inp_files.change(update_visibility, inp_files, timestamps)
+
 demo.launch()