Debug app.py

app.py

@@ -1115,10 +1115,10 @@ def predict(img):
     vflip_tta = VerticalFlip()
 
     img_name = img.name
-    # if img_name.endswith('.tif') or img_name.endswith('.tiff'):
-        # img_data = tif.imread(img_name)
-    # else:
-        # img_data = io.imread(img_name)
+    if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+        origin_img = tif.imread(img_name)
+    else:
+        origin_img = io.imread(img_name)
 
     img_data = pred_transforms(img_name)
     img_data = img_data.to(device)
@@ -1128,7 +1128,7 @@ def predict(img):
         overlap = 0.5
     else:
         overlap = 0.6
-
+    print("start")
     with torch.no_grad():
         img0 = img_data
         outputs0 = sliding_window_inference(
@@ -1231,14 +1231,15 @@ def predict(img):
             outputs = outputs0
 
         pred_mask = post_process(outputs.squeeze(0).cpu().numpy(), device)
-
+    print("prediction end & file write")
     file_path = os.path.join(
         os.getcwd(), img_name.split(".")[0] + "_label.tiff"
     )
 
     tif.imwrite(file_path, pred_mask, compression="zlib")
+    print(np.max(pred_mask))
     # return img_data, seg_rgb, join(os.getcwd(), 'segmentation.tiff')
-    return img_data, pred_mask, file_path
+    return origin_img, pred_mask, file_path
 
 demo = gr.Interface(
     predict,