update model code

app.py

@@ -258,8 +258,22 @@ def predict():
     if (is_mri_image(temp_path)):
         return jsonify({"message": "Not an mri image", "confidence": 0.95, "saved_path": image_save_path})
     a, b = model.check_file(temp_path)
+
+    class ResNetRegression(nn.Module):
+        def __init__(self):
+            super(ResNetRegression, self).__init__()
+            self.model = models.resnet34(pretrained=True)
+            in_features = self.model.fc.in_features
+            # Change output layer for regression
+            self.model.fc = nn.Linear(in_features, 1)
+
+        def forward(self, x):
+            return self.model(x)
+
+    # Initialize Model, Loss, and Optimizer
+    model_new = ResNetRegression()
     image = Image.open(temp_path).convert("RGB")
-    output = model(transform(image).unsqueeze(0))
+    output = model_new(transform(image).unsqueeze(0))
     stage = output.item()
     if stage <= 2.0:
         stage = "Mild"