Update app.py

app.py

@@ -1,7 +1,222 @@
+import os
+from huggingface_hub import hf_hub_download
+config_path=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-sen1floods11", filename="sen1floods11_Prithvi_100M.py", token=os.environ.get("token"))
+ckpt=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-sen1floods11", filename='sen1floods11_Prithvi_100M.pth', token=os.environ.get("token"))
+##########
+
+
+import argparse
+from mmcv import Config
+
+from mmseg.models import build_segmentor
+
+from mmseg.datasets.pipelines import Compose, LoadImageFromFile
+
+import rasterio
+import torch
+
+from mmseg.apis import init_segmentor
+
+from mmcv.parallel import collate, scatter
+
+import numpy as np
+import glob
+import os
+
+import time
+
+import numpy as np
 import gradio as gr
+from functools import partial
+
+import pdb
+
+import matplotlib.pyplot as plt
+
+from skimage import exposure 
+
+def stretch_rgb(rgb):
+    
+    ls_pct=1
+    pLow, pHigh = np.percentile(rgb[~np.isnan(rgb)], (ls_pct,100-ls_pct))
+    img_rescale = exposure.rescale_intensity(rgb, in_range=(pLow,pHigh))
+    
+    return img_rescale
+
+
+def open_tiff(fname):
+    
+    with rasterio.open(fname, "r") as src:
+        
+        data = src.read()
+        
+    return data
+
+def write_tiff(img_wrt, filename, metadata):
+
+    """
+    It writes a raster image to file.
+    :param img_wrt: numpy array containing the data (can be 2D for single band or 3D for multiple bands)
+    :param filename: file path to the output file
+    :param metadata: metadata to use to write the raster to disk
+    :return:
+    """
+
+    with rasterio.open(filename, "w", **metadata) as dest:
+
+        if len(img_wrt.shape) == 2:
+            
+            img_wrt = img_wrt[None]
+
+        for i in range(img_wrt.shape[0]):
+            dest.write(img_wrt[i, :, :], i + 1)
+    
+    return filename
+            
+
+def get_meta(fname):
+    
+    with rasterio.open(fname, "r") as src:
+        
+        meta = src.meta
+        
+    return meta
+
+def preprocess_example(example_list):
+    
+    example_list = [os.path.join(os.path.abspath(''), x) for x in example_list]
+    
+    return example_list
+
+
+def inference_segmentor(model, imgs, custom_test_pipeline=None):
+    """Inference image(s) with the segmentor.
+    Args:
+        model (nn.Module): The loaded segmentor.
+        imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
+            images.
+    Returns:
+        (list[Tensor]): The segmentation result.
+    """
+    cfg = model.cfg
+    device = next(model.parameters()).device  # model device
+    # build the data pipeline
+    test_pipeline = [LoadImageFromFile()] + cfg.data.test.pipeline[1:] if custom_test_pipeline == None else custom_test_pipeline
+    test_pipeline = Compose(test_pipeline)
+    # prepare data
+    data = []
+    imgs = imgs if isinstance(imgs, list) else [imgs]
+    for img in imgs:
+        img_data = {'img_info': {'filename': img}}
+        img_data = test_pipeline(img_data)
+        data.append(img_data)
+    # print(data.shape)
+    
+    data = collate(data, samples_per_gpu=len(imgs))
+    if next(model.parameters()).is_cuda:
+        # data = collate(data, samples_per_gpu=len(imgs))
+        # scatter to specified GPU
+        data = scatter(data, [device])[0]
+    else:
+        # img_metas = scatter(data['img_metas'],'cpu')
+        # data['img_metas'] = [i.data[0] for i in data['img_metas']]
+        
+        img_metas = data['img_metas'].data[0]
+        img = data['img']
+        data = {'img': img, 'img_metas':img_metas}
+    
+    with torch.no_grad():
+        result = model(return_loss=False, rescale=True, **data)
+    return result
+
+
+def inference_on_file(target_image, model, custom_test_pipeline):
+
+    target_image = target_image.name
+
+    time_taken=-1
+    
+    st = time.time()
+    print('Running inference...')
+    result = inference_segmentor(model, target_image, custom_test_pipeline)
+        
+    print("Output has shape: " + str(result[0].shape))
+
+    ##### prep outputs
+    mask = open_tiff(target_image)
+    rgb = stretch_rgb((mask[[3, 2, 1], :, :].transpose((1,2,0))/10000*255).astype(np.uint8))
+    meta = get_meta(target_image)
+    mask = np.where(mask == meta['nodata'], 1, 0)
+    mask = np.max(mask, axis=0)[None]
+    rgb = np.where(mask.transpose((1,2,0)) == 1, 0, rgb)
+    rgb = np.where(rgb < 0, 0, rgb)
+    rgb = np.where(rgb > 255, 255, rgb)
+    
+    prediction = np.where(mask == 1, 0, result[0]*255)
+    et = time.time()
+    time_taken = np.round(et - st, 1)
+    print(f'Inference completed in {str(time_taken)} seconds')
+        
+    return rgb, prediction[0]
+
+def process_test_pipeline(custom_test_pipeline, bands=None):
+    
+    # change extracted bands if necessary
+    if bands is not None:
+        
+        extract_index = [i for i, x in enumerate(custom_test_pipeline) if x['type'] == 'BandsExtract' ]
+        
+        if len(extract_index) > 0:
+            
+            custom_test_pipeline[extract_index[0]]['bands'] = eval(bands)
+            
+    collect_index = [i for i, x in enumerate(custom_test_pipeline) if x['type'].find('Collect') > -1]
+    
+    # adapt collected keys if necessary
+    if len(collect_index) > 0:
+        
+        keys = ['img_info', 'filename', 'ori_filename', 'img', 'img_shape', 'ori_shape', 'pad_shape', 'scale_factor', 'img_norm_cfg']
+        custom_test_pipeline[collect_index[0]]['meta_keys'] = keys
+    
+    return custom_test_pipeline
+
+config = Config.fromfile(config_path)
+config.model.backbone.pretrained=None
+model = init_segmentor(config, ckpt, device='cpu')
+custom_test_pipeline=process_test_pipeline(model.cfg.data.test.pipeline, None)
+
+func = partial(inference_on_file, model=model, custom_test_pipeline=custom_test_pipeline)
 
-def greet(name):
-    return "Hello " + name + "!!"
+with gr.Blocks() as demo:
+   
+    gr.Markdown(value='# Prithvi sen1floods11')
+    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. This demo showcases how the model was finetuned to detect water at a higher resolution than it was trained on (i.e. 10m versus 30m) using Sentinel 2 imagery from on the [sen1floods11 dataset](https://github.com/cloudtostreet/Sen1Floods11). More detailes can be found [here](https://huggingface.co/ibm-nasa-geospatial/Prithvi-100M-sen1floods11).\n
+    The user needs to provide a Sentinel 2 image with all the 12 bands (in the usual Sentinel 2) order in reflectance units multiplied by 10,000 (e.g. to save on space), with the code that is going to pull up Blue, Green, Red, Narrow NIR, SWIR, SWIR 2.
+    ''')
+    with gr.Row():
+        with gr.Column():
+            inp = gr.File()
+            btn = gr.Button("Submit")
+    
+    with gr.Row():
+        gr.Markdown(value='### Input RGB')
+        gr.Markdown(value='### Model prediction (Black: Land; White: Water)')
+        
+    with gr.Row():
+        out1=gr.Image(image_mode='RGB')
+        out2 = gr.Image(image_mode='L')
+    
+    btn.click(fn=func, inputs=inp, outputs=[out1, out2])
+    
+    with gr.Row():
+        gr.Examples(examples=["India_900498_S2Hand.tif",
+                              "Spain_7370579_S2Hand.tif",
+                              "USA_430764_S2Hand.tif"],
+                    inputs=inp,
+                    outputs=[out1, out2],
+                    preprocess=preprocess_example,
+                    fn=func,
+                    cache_examples=True,
+    )
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch(share=True)
+demo.launch()