update model code

app.py

@@ -12,25 +12,28 @@ import os
 import requests
 import requests
 import cloudinary
-import model
+# import model
 import cloudinary.uploader
 from a import main
 import numpy as np
+import torchvision.transforms as transforms
 # Initialize Flask app
 app = Flask(__name__)
 
 GDRIVE_MODEL_URL = "https://drive.google.com/uc?id=1fzKneepaRt_--dzamTcDBM-9d3_dLX7z"
 LOCAL_MODEL_PATH = "checkpoint32.pth"
-
-print(GDRIVE_MODEL_URL)
+d = "https://drive.google.com/uc?id=1GfrlFNoa7E4liMHyMuF73nA21yT9SNSb"
 
 
 def download_file_from_google_drive():
     gdown.download(GDRIVE_MODEL_URL, LOCAL_MODEL_PATH, quiet=False)
 
 
-file_id = "1fzKneepaRt_--dzamTcDBM-9d3_dLX7z"
-destination = "checkpoint32.pth"
+da = "a.pth"
+
+
+def download_file_from_google_drived():
+    gdown.download(d, da, quiet=False)
 
 
 def download_model():
@@ -46,6 +49,7 @@ def download_model():
 
 
 download_file_from_google_drive()
+download_file_from_google_drived()
 
 
 @app.route("/")
@@ -227,6 +231,10 @@ def predict():
     # Save file temporarily
     temp_path = os.path.join(tempfile.gettempdir(), file.filename)
     file.save(temp_path)
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+    ])
     if file.filename.lower().endswith((".png", ".jpg", ".jpeg")):
         image = Image.open(temp_path)
         image_save_path = os.path.join(
@@ -250,7 +258,16 @@ 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)
-    return jsonify({"message": a, "confidence": b, "saved_path": image_save_path})
+    image = Image.open(temp_path).convert("RGB")
+    output = model(transform(image).unsqueeze(0).to(device))
+    stage = output.item()
+    if stage <= 2.0:
+        stage = "Mild"
+    elif stage >= 2.0 and stage <= 3.2:
+        stage = "Moderate"
+    else:
+        stage = "Severe"
+    return jsonify({"message": a, "confidence": b, "stage": stage, "saved_path": image_save_path})
 
 
 if __name__ == "__main__":

model.py

@@ -161,4 +161,4 @@ def check_file(image_path):
     model.to(device)
     model = nn.DataParallel(model)
     output, confidence = test_model(model, test_loader, device)
-    return "control" if output.item() == 0 else "ms", confidence.item()
+    return "No ms detected" if output.item() == 0 else "MS Detected", confidence.item()