Mod 0

app/main.py

@@ -1,19 +1,41 @@
-import numpy as np
 from PIL import Image
+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("/generate")
 def display(text: str):
     return {"yoy":text}
 
-"""@app.post("/predict")
+@app.post("/predict")
 def predict(file: UploadFile):
     img=Image.open(file.file).convert("RGB")
     img=img.resize(480,480)
     img=np.array(img)
     model= get_alzheimer_model()
     label,probability=classify_img(model,img)
-    return {"label":label.item(),"probability":probability.item()}"""
+    return {"label":label.item(),"probability":probability.item()}

app/utils.py

@@ -1,24 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-from torchvision.models import efficientnet_b0
-import torchvision.transforms.functional as tf
-
-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