Delete app.py

######### pull files
import os
from huggingface_hub import hf_hub_download
config_path=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification",
filename="multi_temporal_crop_classification_Prithvi_100M.py",
token=os.environ.get("token"))
ckpt=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification",
filename='multi_temporal_crop_classification_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

cdl_color_map = [{'value': 1, 'label': 'Natural vegetation', 'rgb': (233,255,190)},
{'value': 2, 'label': 'Forest', 'rgb': (149,206,147)},
{'value': 3, 'label': 'Corn', 'rgb': (255,212,0)},
{'value': 4, 'label': 'Soybeans', 'rgb': (38,115,0)},
{'value': 5, 'label': 'Wetlands', 'rgb': (128,179,179)},
{'value': 6, 'label': 'Developed/Barren', 'rgb': (156,156,156)},
{'value': 7, 'label': 'Open Water', 'rgb': (77,112,163)},
{'value': 8, 'label': 'Winter Wheat', 'rgb': (168,112,0)},
{'value': 9, 'label': 'Alfalfa', 'rgb': (255,168,227)},
{'value': 10, 'label': 'Fallow/Idle cropland', 'rgb': (191,191,122)},
{'value': 11, 'label': 'Cotton', 'rgb':(255,38,38)},
{'value': 12, 'label': 'Sorghum', 'rgb':(255,158,15)},
{'value': 13, 'label': 'Other', 'rgb':(0,175,77)}]


def apply_color_map(rgb, color_map=cdl_color_map):


rgb_mapped = rgb.copy()

for map_tmp in cdl_color_map:

for i in range(3):
rgb_mapped[i] = np.where((rgb[0] == map_tmp['value']) & (rgb[1] == map_tmp['value']) & (rgb[2] == map_tmp['value']), map_tmp['rgb'][i], rgb_mapped[i])

return rgb_mapped


def stretch_rgb(rgb):

ls_pct=0
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 process_rgb(input, mask, indexes):


rgb = stretch_rgb((input[indexes, :, :].transpose((1,2,0))/10000*255).astype(np.uint8))
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)

return rgb

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))

##### get metadata mask
input = open_tiff(target_image)
meta = get_meta(target_image)
mask = np.where(input == meta['nodata'], 1, 0)
mask = np.max(mask, axis=0)[None]

rgb1 = process_rgb(input, mask, [2, 1, 0])
rgb2 = process_rgb(input, mask, [8, 7, 6])
rgb3 = process_rgb(input, mask, [14, 13, 12])

result[0] = np.where(mask == 1, 0, result[0])

et = time.time()
time_taken = np.round(et - st, 1)
print(f'Inference completed in {str(time_taken)} seconds')

output=result[0][0] + 1
output = np.vstack([output[None], output[None], output[None]]).astype(np.uint8)
output=apply_color_map(output).transpose((1,2,0))

return rgb1,rgb2,rgb3,output

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)

with gr.Blocks() as demo:

gr.Markdown(value='# Prithvi multi temporal crop classification')
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 classify crop and other land use categories using multi temporal data. More detailes can be found [here](https://huggingface.co/ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification).\n
The user needs to provide an HLS geotiff image, including 18 bands for 3 time-step, and each time-step includes the channels described above (Blue, Green, Red, Narrow NIR, SWIR, SWIR 2) in order.
''')
with gr.Row():
with gr.Column():
inp = gr.File()
btn = gr.Button("Submit")

with gr.Row():
inp1=gr.Image(image_mode='RGB', scale=10, label='T1')
inp2=gr.Image(image_mode='RGB', scale=10, label='T2')
inp3=gr.Image(image_mode='RGB', scale=10, label='T3')
out = gr.Image(image_mode='RGB', scale=10, label='Model prediction')
# gr.Image(value='Legend.png', image_mode='RGB', scale=2, show_label=False)

btn.click(fn=func, inputs=inp, outputs=[inp1, inp2, inp3, out])

with gr.Row():
with gr.Column():
gr.Examples(examples=["chip_102_345_merged.tif",
"chip_104_104_merged.tif",
"chip_109_421_merged.tif"],
inputs=inp,
outputs=[inp1, inp2, inp3, out],
preprocess=preprocess_example,
fn=func,
cache_examples=True)
with gr.Column():
gr.Markdown(value='### Model prediction legend')
gr.Image(value='Legend.png', image_mode='RGB', show_label=False)


demo.launch()

app.py

@@ -1,270 +0,0 @@
-######### pull files
-import os
-from huggingface_hub import hf_hub_download
-config_path=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification", 
-                            filename="multi_temporal_crop_classification_Prithvi_100M.py", 
-                            token=os.environ.get("token"))
-ckpt=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification", 
-                     filename='multi_temporal_crop_classification_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 
-
-cdl_color_map = [{'value': 1, 'label': 'Natural vegetation', 'rgb': (233,255,190)},
-                 {'value': 2, 'label': 'Forest', 'rgb': (149,206,147)},
-                 {'value': 3, 'label': 'Corn', 'rgb': (255,212,0)},
-                 {'value': 4, 'label': 'Soybeans', 'rgb': (38,115,0)},
-                 {'value': 5, 'label': 'Wetlands', 'rgb': (128,179,179)},
-                 {'value': 6, 'label': 'Developed/Barren', 'rgb': (156,156,156)},
-                 {'value': 7, 'label': 'Open Water', 'rgb': (77,112,163)},
-                 {'value': 8, 'label': 'Winter Wheat', 'rgb': (168,112,0)},
-                 {'value': 9, 'label': 'Alfalfa', 'rgb': (255,168,227)},
-                 {'value': 10, 'label': 'Fallow/Idle cropland', 'rgb': (191,191,122)},
-                 {'value': 11, 'label': 'Cotton', 'rgb':(255,38,38)},
-                 {'value': 12, 'label': 'Sorghum', 'rgb':(255,158,15)},
-                 {'value': 13, 'label': 'Other', 'rgb':(0,175,77)}]
-
-
-def apply_color_map(rgb, color_map=cdl_color_map):
-    
-    
-    rgb_mapped = rgb.copy()
-    
-    for map_tmp in cdl_color_map:
-        
-        for i in range(3):
-            rgb_mapped[i] = np.where((rgb[0] == map_tmp['value']) & (rgb[1] == map_tmp['value']) & (rgb[2] == map_tmp['value']), map_tmp['rgb'][i], rgb_mapped[i])
-    
-    return rgb_mapped    
-    
-
-def stretch_rgb(rgb):
-    
-    ls_pct=0
-    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 process_rgb(input, mask, indexes):
-
-    
-    rgb = stretch_rgb((input[indexes, :, :].transpose((1,2,0))/10000*255).astype(np.uint8))
-    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)
-
-    return rgb
-
-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))
-
-    ##### get metadata mask
-    input = open_tiff(target_image)
-    meta = get_meta(target_image)
-    mask = np.where(input == meta['nodata'], 1, 0)
-    mask = np.max(mask, axis=0)[None]
-    
-    rgb1 = process_rgb(input, mask, [2, 1, 0])
-    rgb2 = process_rgb(input, mask, [8, 7, 6])
-    rgb3 = process_rgb(input, mask, [14, 13, 12])
-
-    result[0] = np.where(mask == 1, 0, result[0])
-
-    et = time.time()
-    time_taken = np.round(et - st, 1)
-    print(f'Inference completed in {str(time_taken)} seconds')
-    
-    output=result[0][0] + 1
-    output = np.vstack([output[None], output[None], output[None]]).astype(np.uint8)
-    output=apply_color_map(output).transpose((1,2,0))
-        
-    return rgb1,rgb2,rgb3,output
-
-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)
-
-with gr.Blocks() as demo:
-   
-    gr.Markdown(value='# Prithvi multi temporal crop classification')
-    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 classify crop and other land use categories using multi temporal data. More detailes can be found [here](https://huggingface.co/ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification).\n
-    The user needs to provide an HLS geotiff image, including 18 bands for 3 time-step, and each time-step includes the channels described above (Blue, Green, Red, Narrow NIR, SWIR, SWIR 2) in order.
-    ''')
-    with gr.Row():
-        with gr.Column():
-            inp = gr.File()
-            btn = gr.Button("Submit")
-            
-    with gr.Row():
-        inp1=gr.Image(image_mode='RGB', scale=10, label='T1')
-        inp2=gr.Image(image_mode='RGB', scale=10, label='T2')
-        inp3=gr.Image(image_mode='RGB', scale=10, label='T3')
-        out = gr.Image(image_mode='RGB', scale=10, label='Model prediction')
-        # gr.Image(value='Legend.png', image_mode='RGB', scale=2, show_label=False)
-    
-    btn.click(fn=func, inputs=inp, outputs=[inp1, inp2, inp3, out])
-    
-    with gr.Row():
-        with gr.Column():
-            gr.Examples(examples=["chip_102_345_merged.tif",
-                             "chip_104_104_merged.tif",
-                             "chip_109_421_merged.tif"],
-                            inputs=inp,
-                            outputs=[inp1, inp2, inp3, out],
-                            preprocess=preprocess_example,
-                            fn=func,
-                            cache_examples=True)
-        with gr.Column():
-            gr.Markdown(value='### Model prediction legend')
-            gr.Image(value='Legend.png', image_mode='RGB', show_label=False)
-    
-
-demo.launch()