Update app.py

app.py

@@ -10,6 +10,7 @@ import tempfile
 from pathlib import Path
 from skimage.filters import threshold_otsu
 import torchio as tio
+import psutil
 
 def infer_full_vol(tensor, model):
     tensor = tensor.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, D, H, W] - adding batch and channel dims
@@ -75,6 +76,8 @@ def infer_patch_based(tensor, model, patch_size=64, stride_length=32, stride_wid
             aggregator.add_batch(output, locations)
 
             progress_bar.progress((i + 1) / total_batches)
+            st.text(f"Processing batch {i + 1} of {total_batches}... ({((i + 1) / total_batches) * 100:.2f}% complete)")
+            st.text(f"CPU usage: {psutil.cpu_percent()}% | RAM usage: {psutil.virtual_memory().percent}%")
 
         predicted = aggregator.get_output_tensor().squeeze().numpy()
         
@@ -119,6 +122,15 @@ selected_model = st.selectbox("Select a pretrained model:", model_options)
 mode_options = ["Full volume inference", "Patch-based inference [Default for DS6]"]
 selected_mode = st.selectbox("Select the running mode:", mode_options)
 
+# Parameters for patch-based inference
+if selected_mode == "Patch-based inference [Default for DS6]":
+    patch_size = st.number_input("Patch size:", min_value=1, value=64)
+    stride_length = st.number_input("Stride length:", min_value=1, value=32)
+    stride_width = st.number_input("Stride width:", min_value=1, value=32)
+    stride_depth = st.number_input("Stride depth:", min_value=1, value=16)
+    batch_size = st.number_input("Batch size:", min_value=1, value=10)
+    num_worker = st.number_input("Number of workers:", min_value=1, value=2)
+
 # Process button
 process_button = st.button("Process")
 
@@ -185,7 +197,7 @@ if uploaded_file is not None and process_button:
                     output = infer_full_vol(tensor, model)
                 else:
                     st.info("Running patch-based inference [Default for DS6]...")
-                    output = infer_patch_based(tensor, model)
+                    output = infer_patch_based(tensor, model, patch_size=patch_size, stride_length=stride_length, stride_width=stride_width, stride_depth=stride_depth, batch_size=batch_size, num_worker=num_worker)
             
             st.success("Processing complete.")
             st.write(f"Output tensor shape: `{output.shape}`")
@@ -194,7 +206,7 @@ if uploaded_file is not None and process_button:
                 thresh = threshold_otsu(output)
                 output = output > thresh
             except Exception as error:
-                print(error)
+                st.error(f"Otsu thresholding failed: {error}. Defaulting to a threshold of 0.5.")
                 output = output > 0.5  # exception only if input image seems to have just one color 1.0.
             output = output.astype('uint16')