Update app.py

app.py

@@ -57,7 +57,9 @@ def infer_patch_based(tensor, model, patch_size=64, stride_length=32, stride_wid
                         )
         aggregator = tio.inference.GridAggregator(grid_sampler, overlap_mode="average")
         patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=batch_size, shuffle=False, num_workers=num_worker)
-        for _, patches_batch in enumerate(patch_loader):
+        total_batches = len(patch_loader)
+        progress_bar = st.progress(0)
+        for i, patches_batch in enumerate(patch_loader):
             local_batch = patches_batch['img'][tio.DATA].float()
             local_batch = local_batch / local_batch.max()
             locations = patches_batch[tio.LOCATION]
@@ -72,6 +74,8 @@ def infer_patch_based(tensor, model, patch_size=64, stride_length=32, stride_wid
             output = torch.movedim(output, -3, -1).type(local_batch.type())
             aggregator.add_batch(output, locations)
 
+            progress_bar.progress((i + 1) / total_batches)
+
         predicted = aggregator.get_output_tensor().squeeze().numpy()
         
     return predicted