update app

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.tif* filter=lfs diff=lfs merge=lfs -text

app.py

@@ -5,11 +5,7 @@ config_path=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temp
                             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", 
-<<<<<<< HEAD
                      filename='multi_temporal_crop_classification_Prithvi_100M.pth', 
-=======
-                     filename='multi_temporal_crop_classification_best_mIoU_epoch_66.pth', 
->>>>>>> 889a651 (add files)
                      token=os.environ.get("token"))
 ##########
 import argparse
@@ -40,6 +36,15 @@ import pdb
 
 import matplotlib.pyplot as plt
 
+from skimage import exposure
+
+def stretch_rgb(rgb):
+    
+    ls_pct=0
+    pLow, pHigh = np.percentile(rgb1[~np.isnan(rgb1)], (ls_pct,100-ls_pct))
+    img_rescale = exposure.rescale_intensity(rgb1, in_range=(pLow,pHigh))
+    
+    return img_rescale
 
 def open_tiff(fname):
     
@@ -137,7 +142,6 @@ def inference_on_file(target_image, model, custom_test_pipeline):
 
     # output_image = target_image.replace('.tif', '_pred.tif')
     time_taken=-1
-<<<<<<< HEAD
     st = time.time()
     print('Running inference...')
     result = inference_segmentor(model, target_image, custom_test_pipeline)
@@ -146,9 +150,9 @@ def inference_on_file(target_image, model, custom_test_pipeline):
     ##### get metadata mask
     mask = open_tiff(target_image)
     # rgb = mask[[2, 1, 0], :, :].transpose((1,2,0))
-    rgb1 = mask[[2, 1, 0], :, :].transpose((1,2,0))
-    rgb2 = mask[[8, 7, 6], :, :].transpose((1,2,0))
-    rgb3 = mask[[14, 13, 12], :, :].transpose((1,2,0))
+    rgb1 = stretch_rgb((mask[[2, 1, 0], :, :].transpose((1,2,0))/10000*255).astype(np.unint8))
+    rgb2 = stretch_rgb((mask[[8, 7, 6], :, :].transpose((1,2,0))/10000*255).astype(np.unint8))
+    rgb3 = stretch_rgb((mask[[14, 13, 12], :, :].transpose((1,2,0))/10000*255).astype(np.unint8))
     meta = get_meta(target_image)
     mask = np.where(mask == meta['nodata'], 1, 0)
     mask = np.max(mask, axis=0)[None]
@@ -165,43 +169,11 @@ def inference_on_file(target_image, model, custom_test_pipeline):
     et = time.time()
     time_taken = np.round(et - st, 1)
     print(f'Inference completed in {str(time_taken)} seconds')
+    
+    output = (result[0][0]*18).astype(np.uint8)
+    output = np.vstack([output[None], output[None], output[None]])
         
-    return rgb1,rgb2,rgb3, result[0][0]
-=======
-    try:
-        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
-        mask = open_tiff(target_image)
-        # rgb = mask[[2, 1, 0], :, :].transpose((1,2,0))
-        rgb1 = mask[[2, 1, 0], :, :].transpose((1,2,0))
-        rgb2 = mask[[8, 7, 6], :, :].transpose((1,2,0))
-        rgb3 = mask[[14, 13, 12], :, :].transpose((1,2,0))
-        meta = get_meta(target_image)
-        mask = np.where(mask == meta['nodata'], 1, 0)
-        mask = np.max(mask, axis=0)[None]
-        
-        result[0] = np.where(mask == 1, -1, result[0])
-        
-        ##### Save file to disk
-        meta["count"] = 1
-        meta["dtype"] = "int16"
-        meta["compress"] = "lzw"
-        meta["nodata"] = -1
-        print('Saving output...')
-        # write_tiff(result[0], output_image, meta)
-        et = time.time()
-        time_taken = np.round(et - st, 1)
-        print(f'Inference completed in {str(time_taken)} seconds')
-        
-    except:
-        print(f'Error on image {target_image} \nContinue to next input')
-        
-    return rgb, result[0][0]*255
->>>>>>> 889a651 (add files)
+    return rgb1,rgb2,rgb3,output
 
 def process_test_pipeline(custom_test_pipeline, bands=None):
     
@@ -224,10 +196,6 @@ def process_test_pipeline(custom_test_pipeline, bands=None):
     
     return custom_test_pipeline
 
-<<<<<<< HEAD
-
-=======
->>>>>>> 889a651 (add files)
 config = Config.fromfile(config_path)
 config.model.backbone.pretrained=None
 model = init_segmentor(config, ckpt, device='cpu')
@@ -260,16 +228,6 @@ with gr.Blocks() as demo:
     
     btn.click(fn=func, inputs=inp, outputs=[inp1, inp2, inp3, out])
     
-<<<<<<< HEAD
-    # with gr.Row():
-    #     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.Row():
         gr.Examples(examples=["chip_102_345_merged.tif",
                      "chip_104_104_merged.tif",
@@ -280,6 +238,5 @@ with gr.Blocks() as demo:
                     fn=func,
                     cache_examples=True,
     )
->>>>>>> 889a651 (add files)
 
 demo.launch()