Update app.py

app.py

@@ -113,16 +113,15 @@ def inference(learn, reorder, resample, org_img, input_img, org_size):
             
             # Ensure we have exactly 5 dimensions [1,1,320,320,140]
             if input_img.dim() == 6:
-                print("Removing extra batch dimension")
                 input_img = input_img.squeeze(0)  # Remove extra batch dim
             elif input_img.dim() == 4:
-                print("Adding missing batch dimension")
                 input_img = input_img.unsqueeze(0)  # Add batch dim
-            elif input_img.dim() != 5:
-                raise ValueError(f"Unexpected input dimensions: {input_img.dim()}")
             
-            # Final reshape to exact required dimensions
-            input_img = input_img.view(1, 1, 320, 320, 140)
+            # Verify final shape
+            if input_img.shape != torch.Size([1,1,320,320,140]):
+                input_img = input_img.view(1, 1, 320, 320, 140)
+                print(f"Reshaped to required dimensions: {input_img.shape}")
+            
             print(f"Final input shape to model: {input_img.shape}")
             pred = learn.predict(input_img)
             print(f"Prediction output shape: {pred[0].shape if isinstance(pred, (list, tuple)) else pred.shape}")
@@ -310,6 +309,11 @@ try:
             print(f"Input shape: {tensor_data.shape}")
             print(f"Input dtype: {tensor_data.dtype}")
             
+            # Ensure we have the right dimensions
+            if tensor_data.dim() == 6:
+                tensor_data = tensor_data.squeeze(0)
+                print(f"Removed extra dim - new shape: {tensor_data.shape}")
+            
             # Calculate mean and std
             mean = tensor_data.mean()
             std = tensor_data.std()
@@ -324,8 +328,8 @@ try:
             normalized = (tensor_data - mean) / std
             print(f"Normalized tensor shape: {normalized.shape}")
             
-            # Return as MedImage
-            return MedImage.create(normalized)
+            # Return as plain tensor to avoid MedImage wrapping issues
+            return normalized
             
         except Exception as e:
             print(f"Error in custom normalization: {str(e)}")