fix channels order and no data

app.py

@@ -137,37 +137,36 @@ def inference_segmentor(model, imgs, custom_test_pipeline=None):
 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
+    ##### 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]
-
-    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)
+    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, result[0][0]*255
+        
+    return rgb, prediction[0]
 
 def process_test_pipeline(custom_test_pipeline, bands=None):