update model code

app.py

@@ -299,22 +299,26 @@ def predict():
 
     def remove_module_from_checkpoint(checkpoint):
         new_state_dict = {}
+        print(checkpoint.keys())
         for key, value in checkpoint.items():
             new_key = key.replace("module.", "")
             new_state_dict[new_key] = value
-        checkpoint["model_state_dict"] = new_state_dict
+        checkpoint = new_state_dict
         return checkpoint
     checkpoint = remove_module_from_checkpoint(checkpoint)
     model_new.load_state_dict(checkpoint)
     image = Image.open(temp_path).convert("RGB")
     output = model_new(transform(image).unsqueeze(0))
     stage = output.item()
-    if stage <= 2.0:
-        stage = "Mild"
-    elif stage >= 2.0 and stage <= 3.2:
-        stage = "Moderate"
+    if not a == "No ms detected":
+        if stage <= 2.0:
+            stage = "Mild"
+        elif stage >= 2.0 and stage <= 3.2:
+            stage = "Moderate"
+        else:
+            stage = "Severe"
     else:
-        stage = "Severe"
+        stage = "No ms detected"
     return jsonify({"message": a, "confidence": b, "stage": stage, "saved_path": image_save_path})