fix no data and channels different order

app.py

@@ -163,37 +163,42 @@ def inference_segmentor(model, imgs, custom_test_pipeline=None):
     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
-    # print(type(target_image))
-
-    # output_image = target_image.replace('.tif', '_pred.tif')
     time_taken=-1
     st = time.time()
     print('Running inference...')
-    result = inference_segmentor(model, target_image, custom_test_pipeline)
+    try:
+        result = inference_segmentor(model, target_image, custom_test_pipeline)
+    except:
+        print('Error: Try different channels order.')
+        model.cfg.data.test.pipeline[0]['channels_last'] = True
+        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 = stretch_rgb((mask[[2, 1, 0], :, :].transpose((1,2,0))/10000*255).astype(np.uint8))
-    rgb2 = stretch_rgb((mask[[8, 7, 6], :, :].transpose((1,2,0))/10000*255).astype(np.uint8))
-    rgb3 = stretch_rgb((mask[[14, 13, 12], :, :].transpose((1,2,0))/10000*255).astype(np.uint8))
+    input = open_tiff(target_image)
     meta = get_meta(target_image)
-    mask = np.where(mask == meta['nodata'], 1, 0)
+    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, -1, result[0])
+    result[0] = np.where(mask == 1, 0, 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')