Final Application

.gitignore

@@ -0,0 +1,2 @@
+.venv/
+app/__pycache__/

app/main.py

@@ -1,43 +1,22 @@
 from PIL import Image
+import cv2
 import numpy as np
-import torch
-import torch.nn as nn
-from torchvision.models import efficientnet_b0
-import torchvision.transforms.functional as tf
 from fastapi import FastAPI, UploadFile
+from .utils import classify_img,get_alzheimer_model
 
 app=FastAPI(title="Alzheimer Detection API")
 
-def classify_img(model,img):
-    img=tf.to_tensor(img)
-    img=img.unsqueeze(0)
-    with torch.no_grad():
-        predict=model(img)
-        predict=nn.functional.softmax(predict,1)
-        label=torch.argmax(predict)
-        probability=torch.max(predict)
-        return label,probability
-
-def get_alzheimer_model():
-    model=efficientnet_b0(weights=None)
-    in_features=model.classifier[1].in_features
-    model.classifier[1]=nn.Linear(in_features=in_features,out_features=4)
-    weights=torch.load("alzheimer_weight.pth",map_location="cpu")
-    model.load_state_dict(weights)
-    model.eval()
-    return model
-
-
-
 @app.get("/")
 def display():
     return "Welcome to Alzheimer Detection Api"
 
 @app.post("/predict")
 def predict(file: UploadFile):
-    img=Image.open(file.file).convert("RGB")
-    img=img.resize(480,480)
+    img=Image.open(file.file)
     img=np.array(img)
+    if len(img.shape)==2:
+        img=cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)
+    img=cv2.resize(img,(480,480))
     model= get_alzheimer_model()
     label,probability=classify_img(model,img)
     return {"label":label.item(),"probability":probability.item()}

app/utils.py

@@ -0,0 +1,29 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision
+import os
+from torchvision.models import efficientnet_b0
+import torchvision.transforms.functional as tf
+
+root_dir=os.path.dirname(os.path.abspath(__file__))
+weights_pth=os.path.join(root_dir,"alzheimer_weight.pth")
+
+def classify_img(model,img):
+    img=tf.to_tensor(img)
+    img=img.unsqueeze(0)
+    with torch.no_grad():
+        predict=model(img)
+        predict=nn.functional.softmax(predict,1)
+        label=torch.argmax(predict)
+        probability=torch.max(predict)
+        return label,probability
+
+def get_alzheimer_model():
+    model=efficientnet_b0(weights=None)
+    in_features=model.classifier[1].in_features
+    model.classifier[1]=nn.Linear(in_features=in_features,out_features=4)
+    weights=torch.load(weights_pth,map_location="cpu")
+    model.load_state_dict(weights)
+    model.eval()
+    return model

requirements.txt

@@ -3,4 +3,5 @@ torch==2.0.1
 fastapi==0.104.1
 uvicorn==0.24.0
 python-multipart==0.0.6
+opencv-python-headless==4.5.4.60
 torchvision==0.15.2