Made changes to mediseg.py

app/mediseg.py

@@ -117,7 +117,7 @@ class UNetMulti(nn.Module):
 
 def process_brain_tumor(image: Image.Image, model_path=os.path.join("models", "brain_tumor_unet_multiclass.pth")) -> str:
     model = UNetMulti(in_channels=3, out_channels=4).to(device)
-    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
     model.eval()
     
     transform_img = transforms.Compose([
@@ -204,7 +204,7 @@ class UNetBinary(nn.Module):
 
 def process_endoscopy(image: Image.Image, model_path=os.path.join("models", "endoscopy_unet.pth")) -> str:
     model = UNetBinary(in_channels=3, out_channels=1).to(device)
-    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
     model.eval()
     
     transform_img = transforms.Compose([
@@ -284,7 +284,7 @@ def process_pneumonia(image: Image.Image, model_path=os.path.join("models", "pne
     model = models.resnet18(pretrained=False)
     num_ftrs = model.fc.in_features
     model.fc = nn.Linear(num_ftrs, 2)  # 2 classes: normal and pneumonia
-    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
     model.to(device)
     model.eval()