prova full vol

app.py

@@ -10,6 +10,18 @@ import os
 import tempfile
 from pathlib import Path
 import pandas as pd
+from skimage.filters import threshold_otsu
+
+def infer_full_vol(tensor, model):
+    tensor = torch.movedim(tensor, -1, -3)
+    tensor = tensor / tensor.max()
+    with torch.no_grad():
+        output = model(tensor) 
+        if type(output) is tuple or type(output) is list:
+            output = output[0]
+        output = torch.sigmoid(output)
+    output = torch.movedim(output, -3, -1).type(tensor.type())
+    return output.squeeze(0)
 
 # Set page configuration
 st.set_page_config(
@@ -110,14 +122,18 @@ if uploaded_file is not None and process_button:
 
             # Process the tensor through the model
             with st.spinner('Processing the tensor through the model...'):
-                with torch.no_grad():
-                    output = model(tensor) 
-                    if isinstance(output, tuple):
-                        output = output[0]
-            output = output.squeeze(0)
+                output = infer_full_vol(tensor, model)
             
             st.success("Processing complete.")
             st.write(f"Output tensor shape: `{output.shape}`")
+            
+            try:
+                thresh = threshold_otsu(output)
+                output = output > thresh
+            except Exception as error:
+                print(error)
+                output = output > 0.5  # exception only if input image seems to have just one color 1.0.
+            output = output.astype('uint16')
 
             # Convert output to NumPy array
             output_np = output.detach().cpu().numpy()

requirements.txt

@@ -3,4 +3,5 @@ torch
 pytorch_lightning
 scipy
 transformers
-torchvision
+torchvision
+scikit-image