update model code

app.py

@@ -232,8 +232,8 @@ def predict():
         image_save_path = os.path.join(
             tempfile.gettempdir(), file.filename.lower())
         image.save(image_save_path)
-
-    return jsonify({"message": model.check_file(temp_path), "saved_path": image_save_path})
+    a, b = model.check_file(temp_path)[0]
+    return jsonify({"message": a, confidence: b, "saved_path": image_save_path})
 
 
 if __name__ == "__main__":

model.py

@@ -142,7 +142,7 @@ def check_file(image_path):
                 # Get confidence score and prediction
                 confidence, d = torch.max(probabilities, 1)
                 print(confidence)
-                return predicted
+                return predicted, confidence
 
     def remove_module_from_checkpoint(checkpoint):
         new_state_dict = {}
@@ -160,6 +160,5 @@ def check_file(image_path):
     model.eval()
     model.to(device)
     model = nn.DataParallel(model)
-    output = test_model(model, test_loader, device)
-    print(output)
-    return "control" if output.item() == 0 else "ms"
+    output, confidence = test_model(model, test_loader, device)
+    return "control" if output.item() == 0 else "ms", confidence.item()