fix no data issue

app.py

@@ -137,28 +137,23 @@ def inference_on_file(target_image, model, custom_test_pipeline):
     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 = mask[[2, 1, 0], :, :].transpose((1,2,0))
     rgb = mask[[5, 3, 2], :, :].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)
+    rgb = np.where(mask.transpose((1,2,0)) == 1, 0, rgb)
+    rgb = np.where(rgb < 0, 0, rgb)
+    rgb = np.where(rgb > 1, 1, 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):