fullvol auto-padding added

app.py

@@ -1,8 +1,10 @@
 import streamlit as st
 import json
+import math
 import numpy as np
 import nibabel as nib
 import torch
+import torch.nn.functional as F
 import scipy.io
 from io import BytesIO
 from transformers import AutoModel
@@ -15,11 +17,32 @@ from skimage.filters import threshold_otsu
 def infer_full_vol(tensor, model):
     tensor = torch.movedim(tensor, -1, -3)
     tensor = tensor / tensor.max()
+    
+    sizes = tensor.shape[-3:]
+    new_sizes = [math.ceil(s / 16) * 16 for s in sizes]
+    total_pads = [new_size - s for s, new_size in zip(sizes, new_sizes)]
+    pad_before = [pad // 2 for pad in total_pads]
+    pad_after = [pad - pad_before[i] for i, pad in enumerate(total_pads)]
+    padding = []
+    for i in reversed(range(len(pad_before))):
+        padding.extend([pad_before[i], pad_after[i]])
+    tensor = F.pad(tensor, padding)
+    
     with torch.no_grad():
         output = model(tensor) 
         if type(output) is tuple or type(output) is list:
             output = output[0]
         output = torch.sigmoid(output)
+        
+    slices = [slice(None)] * output.dim()
+    for i in range(len(pad_before)):
+        dim = -3 + i
+        start = pad_before[i]
+        size = sizes[i]
+        end = start + size
+        slices[dim] = slice(start, end)
+    output = output[tuple(slices)]
+    
     output = torch.movedim(output, -3, -1).type(tensor.type())
     return output.squeeze().detach().cpu().numpy()