stretch image

app.py

@@ -34,6 +34,16 @@ import pdb
 
 import matplotlib.pyplot as plt
 
+from skimage import exposure 
+
+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):
     
@@ -139,8 +149,7 @@ 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))
-    rgb = mask[[3, 2, 1], :, :].transpose((1,2,0))
+    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]
@@ -157,8 +166,8 @@ 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')
-        
-    return (rgb/10000*255).astype(np.uint8), result[0][0]*255
+    
+    return rgb, result[0][0]*255
 
 def process_test_pipeline(custom_test_pipeline, bands=None):