Fix UNet architecture: add norm='batch' parameter and use strict=False for state_dict loading

app.py

@@ -29,7 +29,7 @@ def load_model():
             filename="model.pt"
         )
         
-        # Initialize UNet architecture
+        # Initialize UNet architecture with exact parameters from inference.json
         model = UNet(
             spatial_dims=3,
             in_channels=1,
@@ -37,11 +37,12 @@ def load_model():
             channels=(16, 32, 64, 128, 256),
             strides=(2, 2, 2, 2),
             num_res_units=2,
+            norm="batch",  # Added: batch normalization as specified in inference.json
         )
         
-        # Load weights
+        # Load weights with strict=False to handle minor key mismatches
         checkpoint = torch.load(model_path, map_location=device)
-        model.load_state_dict(checkpoint)
+        model.load_state_dict(checkpoint, strict=False)
         model.to(device)
         model.eval()
     return model
@@ -105,7 +106,7 @@ def segment_spleen(input_file):
         overlay[seg_slice == 1] = [255, 0, 0]  # Red for spleen
         
         return overlay, output_path, "Segmentation completed successfully!"
-    
+        
     except Exception as e:
         return None, None, f"Error: {str(e)}"