Update app

app.py

@@ -1,215 +1,28 @@
-#### pull files from hub
-from huggingface_hub import hf_hub_download
-import os
-yaml_file_path=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M", filename="Prithvi_100M_config.yaml", token=os.environ.get("token"))
-checkpoint=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M", filename='Prithvi_100M.pt', token=os.environ.get("token"))
-model_def=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M", filename='Prithvi.py', token=os.environ.get("token"))
-os.system(f'cp {model_def} .')
-#####
-import argparse
-import functools
-import os
-from typing import List
 
-import numpy as np
-import rasterio
+import os
 import torch
 import yaml
-from einops import rearrange
-
-from Prithvi import MaskedAutoencoderViT
+import numpy as np
 import gradio as gr
+from einops import rearrange
 from functools import partial
+from huggingface_hub import hf_hub_download
 
+# pull files from hub
+token = os.environ.get("HF_TOKEN", None)
+config_path = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-1.0-100M",
+                                 filename="config.json", token=token)
+checkpoint = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-1.0-100M",
+                             filename='Prithvi_EO_V1_100M.pt', token=token)
+model_def = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-1.0-100M",
+                            filename='prithvi_mae.py', token=token)
+model_inference = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-1.0-100M",
+                                  filename='inference.py', token=token)
+os.system(f'cp {model_def} .')
+os.system(f'cp {model_inference} .')
 
-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
-
-    return img, meta
-
-
-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])
-
+from prithvi_mae import PrithviMAE
+from inference import process_channel_group, _convert_np_uint8, load_example, run_model
 
 def extract_rgb_imgs(input_img, rec_img, mask_img, channels, mean, std):
     """ Wrapper function to save Geotiff images (original, reconstructed, masked) per timestamp.
@@ -230,24 +43,31 @@ def extract_rgb_imgs(input_img, rec_img, mask_img, channels, mean, std):
     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)
+                                                   channels=channels,
+                                                   mean=mean,
+                                                   std=std)
 
         rgb_mask = mask_img[channels, t, :, :] * rgb_orig
 
         # extract images
-        rgb_orig_list.append(_convert_np_uint8(rgb_orig))
-        rgb_mask_list.append(_convert_np_uint8(rgb_mask))
-        rgb_pred_list.append(_convert_np_uint8(rgb_pred))
-        
+        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 = 3 - 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, config_path: str, checkpoint: str, mask_ratio: float = None):
     try:
         data_files = [x.name for x in data_files]
         print('Path extracted from example')
@@ -257,24 +77,18 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
     # Get parameters --------
     print('This is the printout', data_files)
 
-    with open(yaml_file_path, 'r') as f:
-        params = yaml.safe_load(f)
-
-    model_params = params["model_args"]
-    # data related
-    train_params = params["train_params"]
-    num_frames = model_params['num_frames']
-    img_size = model_params['img_size']
-    bands = train_params['bands']
-    mean = train_params['data_mean']
-    std = train_params['data_std']
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)['pretrained_cfg']
 
     batch_size = 8
+    bands = config['bands']
+    num_frames = len(data_files)
+    mean = config['mean']
+    std = config['std']
+    img_size = config['img_size']
+    mask_ratio = mask_ratio or config['mask_ratio']
 
-    mask_ratio = train_params['mask_ratio'] if mask_ratio is None else mask_ratio
-
-    # We must have *num_frames* files to build one example!
-    assert len(data_files) == num_frames, "File list must be equal to expected number of frames."
+    assert num_frames <= 3, "Demo only supports up to three timestamps"
 
     if torch.cuda.is_available():
         device = torch.device('cuda')
@@ -289,16 +103,23 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
 
     # Create model and load checkpoint -------------------------------------------------------------
 
-    model = MaskedAutoencoderViT(
-            **model_params)
+    config.update(
+        num_frames=num_frames,
+    )
+
+    model = PrithviMAE(**config)
 
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
     print(f"\n--> Model has {total_params:,} parameters.\n")
 
     model.to(device)
 
-    state_dict = torch.load(checkpoint, map_location=device)
-    model.load_state_dict(state_dict)
+    state_dict = torch.load(checkpoint, map_location=device, weights_only=False)
+    # discard fixed pos_embedding weight
+    for k in list(state_dict.keys()):
+        if 'pos_embed' in k:
+            del state_dict[k]
+    model.load_state_dict(state_dict, strict=False)
     print(f"Loaded checkpoint from {checkpoint}")
 
     # Running model --------------------------------------------------------------------------------
@@ -348,37 +169,35 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
     for d in meta_data:
         d.update(count=3, dtype='uint8', compress='lzw', nodata=0)
 
-    # save_rgb_imgs(batch[0, ...], rec_imgs_full[0, ...], mask_imgs_full[0, ...],
-    #               channels, mean, std, output_dir, meta_data)
-
     outputs = extract_rgb_imgs(batch_full[0, ...], rec_imgs_full[0, ...], mask_imgs_full[0, ...],
                   channels, mean, std)
 
-
     print("Done!")
 
     return outputs
 
 
+run_inference = partial(predict_on_images, config_path=config_path,checkpoint=checkpoint)
 
+with gr.Blocks() as demo:
 
-func = partial(predict_on_images, yaml_file_path=yaml_file_path,checkpoint=checkpoint)
+    gr.Markdown(value='# Prithvi-EO-1.0 image reconstruction demo')
+    gr.Markdown(value='''
+Check out our newest model: [Prithvi-EO-2.0-Demo](https://huggingface.co/spaces/ibm-nasa-geospatial/Prithvi-EO-2.0-Demo).
 
-def preprocess_example(example_list):
-    print('######## preprocessing here ##########')
-    example_list = [os.path.join(os.path.abspath(''), x) for x in example_list]
-    
-    return example_list
-    
+Prithvi is a first-of-its-kind temporal Vision transformer pretrained by the IBM and NASA team on continental US Harmonised Landsat Sentinel 2 (HLS) data. 
+Particularly, the model adopts a self-supervised encoder developed with a ViT architecture and Masked AutoEncoder learning strategy, with a MSE as a loss function. 
+The model includes spatial attention across multiple patchies and also temporal attention for each patch. 
+More info about the model and its weights are available [here](https://huggingface.co/ibm-nasa-geospatial/Prithvi-100M).\n
 
-with gr.Blocks() as demo:
-    
-    gr.Markdown(value='# Prithvi image reconstruction demo')
-    gr.Markdown(value='''Prithvi is a first-of-its-kind temporal Vision transformer pretrained by the IBM and NASA team on continental US Harmonised Landsat Sentinel 2 (HLS) data. Particularly, the model adopts a self-supervised encoder developed with a ViT architecture and Masked AutoEncoder learning strategy, with a MSE as a loss function. The model includes spatial attention across multiple patchies and also temporal attention for each patch. More info about the model and its weights are available [here](https://huggingface.co/ibm-nasa-geospatial/Prithvi-100M).\n
-This demo showcases the image reconstracting over three timestamps, with the user providing a set of three HLS images and the model randomly masking out some proportion of the images and then reconstructing them based on the not masked portion of the images.\n
-The user needs to provide three HLS geotiff images, including the following channels in reflectance units: Blue, Green, Red, Narrow NIR, SWIR, SWIR 2.
+This demo showcases the image reconstruction over one to three timestamps. 
+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
 
-Check out our newest model: [Prithvi-EO-2.0-Demo](https://huggingface.co/spaces/ibm-nasa-geospatial/Prithvi-EO-2.0-Demo).
+The user needs to provide the HLS geotiff images, including the following channels in reflectance units: Blue, Green, Red, Narrow NIR, SWIR, SWIR 2. 
+Some example images are provided at the end of this page.
 ''')
     with gr.Row():
         with gr.Column():
@@ -386,73 +205,55 @@ Check out our newest model: [Prithvi-EO-2.0-Demo](https://huggingface.co/spaces/
             # inp_slider = gr.Slider(0, 100, value=50, label="Mask ratio", info="Choose ratio of masking between 0 and 100", elem_id='slider'),
             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='## Input time series')
         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='## Reconstructed images*')
+
+    original = []
+    masked = []
+    predicted = []
+    timestamps = []
+    for t in range(3):
+        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', show_label=False, show_fullscreen_button=False))
+                masked.append(gr.Image(image_mode='RGB', show_label=False, show_fullscreen_button=False))
+                predicted.append(gr.Image(image_mode='RGB', show_label=False, show_fullscreen_button=False))
+
+    gr.Markdown(value='\* The reconstructed images include the ground truth unmasked patches.')
+
+    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__), "HLS.L30.T13REN.2018013T172747.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
-                      os.path.join(os.path.dirname(__file__), "HLS.L30.T13REN.2018029T172738.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
-                      os.path.join(os.path.dirname(__file__), "HLS.L30.T13REN.2018061T172724.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif")]],
-                     [[os.path.join(os.path.dirname(__file__), "HLS.L30.T17RMP.2018004T155509.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
-                      os.path.join(os.path.dirname(__file__), "HLS.L30.T17RMP.2018036T155452.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
-                      os.path.join(os.path.dirname(__file__), "HLS.L30.T17RMP.2018068T155438.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif")]],
-                     [[os.path.join(os.path.dirname(__file__), "HLS.L30.T18TVL.2018029T154533.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
-                      os.path.join(os.path.dirname(__file__), "HLS.L30.T18TVL.2018141T154435.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
-                      os.path.join(os.path.dirname(__file__), "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],
-                    # preprocess=preprocess_example,
-                    fn=func,
-                    cache_examples=True
+        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"),
+            os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T13REN.2018029T172738.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
+            os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T13REN.2018061T172724.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif")
+        ]],[[
+            os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T17RMP.2018004T155509.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
+            os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T17RMP.2018036T155452.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
+            os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T17RMP.2018068T155438.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif")
+        ]],[[
+            os.path.join(os.path.dirname(__file__), "examples/HLS.L30.T18TVL.2018029T154533.v2.0.B02.B03.B04.B05.B06.B07_cropped.tif"),
+            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=original + masked + predicted,
+        fn=run_inference,
+        cache_examples=True
     )
-    
 
-demo.launch()
+    def update_visibility(files):
+        timestamps = [gr.Column(visible=t < len(files)) for t in range(3)]
+
+        return timestamps
+
+    inp_files.change(update_visibility, inp_files, timestamps)
+
+demo.launch() # share=True, ssr_mode=False