Update app.py

app.py

@@ -1,10 +1,11 @@
 # =========================================
 # IMPORTS
 # =========================================
-import io
+import os
 import cv2
 import nibabel as nib
 import numpy as np
+
 import torch
 import torch.nn as nn
 import torchvision.models as models
@@ -33,42 +34,67 @@ print("Using device:", DEVICE)
 # FASTAPI
 # =========================================
 app = FastAPI(
-    title="Parkinsons DATSCAN Ensemble API",
+    title="Parkinson DATSCAN Ensemble API",
     version="1.0"
 )
 
 # =========================================
-# PREPROCESSING
+# LOAD NIFTI
 # =========================================
 def load_nifti_from_bytes(file_bytes):
-    temp_path = "temp.nii"
+
+    temp_path = "temp_upload.nii"
 
     with open(temp_path, "wb") as f:
         f.write(file_bytes)
 
     volume = nib.load(temp_path).get_fdata()
+
     volume = np.squeeze(volume)
 
     return volume
 
-
+# =========================================
+# PREPROCESS 2D
+# =========================================
 def preprocess_2d(volume):
+
     depth = volume.shape[2]
 
-    idx1 = np.linspace(0, depth//3 - 1, 10).astype(int)
-    idx2 = np.linspace(depth//3, 2*depth//3 - 1, 10).astype(int)
-    idx3 = np.linspace(2*depth//3, depth - 1, 12).astype(int)
+    idx1 = np.linspace(
+        0,
+        depth // 3 - 1,
+        10
+    ).astype(int)
+
+    idx2 = np.linspace(
+        depth // 3,
+        2 * depth // 3 - 1,
+        10
+    ).astype(int)
+
+    idx3 = np.linspace(
+        2 * depth // 3,
+        depth - 1,
+        12
+    ).astype(int)
 
     def make_channel(indices):
+
         slices = []
 
         for i in indices:
+
             sl = volume[:, :, i]
 
             sl = sl - sl.min()
             sl = sl / (sl.max() + 1e-6)
 
-            sl = cv2.resize(sl, (IMG_SIZE, IMG_SIZE))
+            sl = cv2.resize(
+                sl,
+                (IMG_SIZE, IMG_SIZE)
+            )
+
             slices.append(sl)
 
         return np.mean(slices, axis=0)
@@ -79,50 +105,67 @@ def preprocess_2d(volume):
 
     img = np.stack([r, g, b], axis=0)
 
-    return torch.tensor(img, dtype=torch.float32).unsqueeze(0)
+    tensor = torch.tensor(
+        img,
+        dtype=torch.float32
+    )
+
+    tensor = tensor.unsqueeze(0)
 
+    return tensor
 
+# =========================================
+# PREPROCESS 3D
+# =========================================
 def preprocess_3d(volume):
+
     depth = volume.shape[2]
 
-    indices = np.linspace(0, depth - 1, NUM_SLICES).astype(int)
+    indices = np.linspace(
+        0,
+        depth - 1,
+        NUM_SLICES
+    ).astype(int)
 
     slices = []
 
     for i in indices:
+
         sl = volume[:, :, i]
 
         sl = sl - sl.min()
         sl = sl / (sl.max() + 1e-6)
 
-        sl = cv2.resize(sl, (IMG_SIZE, IMG_SIZE))
+        sl = cv2.resize(
+            sl,
+            (IMG_SIZE, IMG_SIZE)
+        )
+
         slices.append(sl)
 
     vol = np.stack(slices)
 
-    vol = torch.tensor(vol, dtype=torch.float32)
-
-    vol = vol.unsqueeze(0).unsqueeze(0)
+    tensor = torch.tensor(
+        vol,
+        dtype=torch.float32
+    )
 
-    return vol
+    tensor = tensor.unsqueeze(0)
+    tensor = tensor.unsqueeze(0)
 
+    return tensor
 
 # =========================================
-# DENSENET MODELS
+# DENSENET121
 # =========================================
 class DenseNet121Model(nn.Module):
+
     def __init__(self):
-        super().__init__()
 
-        self.base = models.densenet121(weights=None)
+        super().__init__()
 
-        self.base.features.conv0 = nn.Conv2d(
-            3,
-            64,
-            kernel_size=7,
-            stride=2,
-            padding=3,
-            bias=False
+        self.base = models.densenet121(
+            weights=None
         )
 
         self.base.classifier = nn.Linear(
@@ -131,22 +174,20 @@ class DenseNet121Model(nn.Module):
         )
 
     def forward(self, x):
-        return self.base(x)
 
+        return self.base(x)
 
+# =========================================
+# DENSENET169
+# =========================================
 class DenseNet169Model(nn.Module):
+
     def __init__(self):
-        super().__init__()
 
-        self.base = models.densenet169(weights=None)
+        super().__init__()
 
-        self.base.features.conv0 = nn.Conv2d(
-            3,
-            64,
-            kernel_size=7,
-            stride=2,
-            padding=3,
-            bias=False
+        self.base = models.densenet169(
+            weights=None
         )
 
         self.base.classifier = nn.Linear(
@@ -155,22 +196,20 @@ class DenseNet169Model(nn.Module):
         )
 
     def forward(self, x):
-        return self.base(x)
 
+        return self.base(x)
 
+# =========================================
+# DENSENET201
+# =========================================
 class DenseNet201Model(nn.Module):
+
     def __init__(self):
-        super().__init__()
 
-        self.base = models.densenet201(weights=None)
+        super().__init__()
 
-        self.base.features.conv0 = nn.Conv2d(
-            3,
-            64,
-            kernel_size=7,
-            stride=2,
-            padding=3,
-            bias=False
+        self.base = models.densenet201(
+            weights=None
         )
 
         self.base.classifier = nn.Linear(
@@ -179,188 +218,380 @@ class DenseNet201Model(nn.Module):
         )
 
     def forward(self, x):
-        return self.base(x)
 
+        return self.base(x)
 
 # =========================================
 # 3D CNN
 # =========================================
 class CNN3D(nn.Module):
+
     def __init__(self):
+
         super().__init__()
 
         self.net = nn.Sequential(
 
-            nn.Conv3d(1, 16, 3, padding=1),
+            nn.Conv3d(
+                1,
+                16,
+                kernel_size=3,
+                padding=1
+            ),
+
             nn.ReLU(),
+
             nn.MaxPool3d(2),
 
-            nn.Conv3d(16, 32, 3, padding=1),
+            nn.Conv3d(
+                16,
+                32,
+                kernel_size=3,
+                padding=1
+            ),
+
             nn.ReLU(),
+
             nn.MaxPool3d(2),
 
-            nn.Conv3d(32, 64, 3, padding=1),
+            nn.Conv3d(
+                32,
+                64,
+                kernel_size=3,
+                padding=1
+            ),
+
             nn.ReLU(),
-            nn.MaxPool3d(2)
 
+            nn.MaxPool3d(2)
         )
 
         self.fc = nn.Sequential(
-            nn.Linear(64 * 4 * 16 * 16, 256),
+
+            nn.Linear(
+                64 * 4 * 16 * 16,
+                256
+            ),
+
             nn.ReLU(),
+
             nn.Dropout(0.3),
-            nn.Linear(256, NUM_CLASSES)
+
+            nn.Linear(
+                256,
+                NUM_CLASSES
+            )
         )
 
     def forward(self, x):
+
         x = self.net(x)
-        x = x.view(x.size(0), -1)
-        return self.fc(x)
 
+        x = x.view(
+            x.size(0),
+            -1
+        )
+
+        x = self.fc(x)
+
+        return x
 
 # =========================================
 # LOAD MODELS
 # =========================================
 def load_model(model, path):
+
+    print(f"Loading {path}")
+
+    state = torch.load(
+        path,
+        map_location=DEVICE
+    )
+
     model.load_state_dict(
-        torch.load(path, map_location=DEVICE)
+        state,
+        strict=False
     )
 
     model.to(DEVICE)
+
     model.eval()
 
-    print(f"Loaded: {path}")
+    print(f"Loaded {path}")
 
     return model
 
+# =========================================
+# LOAD ALL
+# =========================================
+model121 = load_model(
+    DenseNet121Model(),
+    "densenet121.pth"
+)
+
+model169 = load_model(
+    DenseNet169Model(),
+    "densenet169.pth"
+)
 
-model121 = load_model(DenseNet121Model(), "densenet121.pth")
-model169 = load_model(DenseNet169Model(), "densenet169.pth")
-model201 = load_model(DenseNet201Model(), "densenet201.pth")
-model3d = load_model(CNN3D(), "cnn3d.pth")
+model201 = load_model(
+    DenseNet201Model(),
+    "densenet201.pth"
+)
 
+model3d = load_model(
+    CNN3D(),
+    "cnn3d.pth"
+)
 
 # =========================================
-# PREDICTION HELPERS
+# SINGLE PREDICTION
 # =========================================
 def predict_model(model, tensor):
+
     tensor = tensor.to(DEVICE)
 
     with torch.no_grad():
+
         out = model(tensor)
 
-        probs = torch.softmax(out, dim=1)[0]
+        probs = torch.softmax(
+            out,
+            dim=1
+        )[0]
 
-    pred_idx = torch.argmax(probs).item()
+    pred_idx = torch.argmax(
+        probs
+    ).item()
+
+    result = {
+
+        "prediction":
+            LABELS[pred_idx],
+
+        "class_id":
+            pred_idx,
+
+        "confidence":
+            round(
+                float(
+                    probs[pred_idx]
+                ) * 100,
+                2
+            ),
 
-    return {
-        "prediction": LABELS[pred_idx],
-        "class_id": pred_idx,
-        "confidence": round(float(probs[pred_idx]) * 100, 2),
         "probabilities": {
-            LABELS[i]: round(float(probs[i]) * 100, 2)
+
+            LABELS[i]:
+                round(
+                    float(
+                        probs[i]
+                    ) * 100,
+                    2
+                )
+
             for i in range(NUM_CLASSES)
         }
-    }, probs
+    }
 
+    return result, probs
 
 # =========================================
 # ROOT
 # =========================================
 @app.get("/")
 def home():
+
     return {
-        "message": "Parkinson DATSCAN Ensemble API Running",
-        "classes": LABELS
-    }
 
+        "message":
+            "Parkinson DATSCAN Ensemble API Running",
+
+        "classes":
+            LABELS
+    }
 
 # =========================================
-# INDIVIDUAL ENDPOINTS
+# DENSENET121 ENDPOINT
 # =========================================
 @app.post("/predict/densenet121")
-async def predict_121(file: UploadFile = File(...)):
+async def predict_121(
+    file: UploadFile = File(...)
+):
+
+    file_bytes = await file.read()
+
+    volume = load_nifti_from_bytes(
+        file_bytes
+    )
 
-    volume = load_nifti_from_bytes(await file.read())
     tensor = preprocess_2d(volume)
 
-    result, _ = predict_model(model121, tensor)
+    result, _ = predict_model(
+        model121,
+        tensor
+    )
 
     return JSONResponse(result)
 
-
+# =========================================
+# DENSENET169 ENDPOINT
+# =========================================
 @app.post("/predict/densenet169")
-async def predict_169(file: UploadFile = File(...)):
+async def predict_169(
+    file: UploadFile = File(...)
+):
+
+    file_bytes = await file.read()
+
+    volume = load_nifti_from_bytes(
+        file_bytes
+    )
 
-    volume = load_nifti_from_bytes(await file.read())
     tensor = preprocess_2d(volume)
 
-    result, _ = predict_model(model169, tensor)
+    result, _ = predict_model(
+        model169,
+        tensor
+    )
 
     return JSONResponse(result)
 
-
+# =========================================
+# DENSENET201 ENDPOINT
+# =========================================
 @app.post("/predict/densenet201")
-async def predict_201(file: UploadFile = File(...)):
+async def predict_201(
+    file: UploadFile = File(...)
+):
+
+    file_bytes = await file.read()
+
+    volume = load_nifti_from_bytes(
+        file_bytes
+    )
 
-    volume = load_nifti_from_bytes(await file.read())
     tensor = preprocess_2d(volume)
 
-    result, _ = predict_model(model201, tensor)
+    result, _ = predict_model(
+        model201,
+        tensor
+    )
 
     return JSONResponse(result)
 
-
+# =========================================
+# 3D CNN ENDPOINT
+# =========================================
 @app.post("/predict/cnn3d")
-async def predict_3d(file: UploadFile = File(...)):
+async def predict_cnn3d(
+    file: UploadFile = File(...)
+):
+
+    file_bytes = await file.read()
+
+    volume = load_nifti_from_bytes(
+        file_bytes
+    )
 
-    volume = load_nifti_from_bytes(await file.read())
     tensor = preprocess_3d(volume)
 
-    result, _ = predict_model(model3d, tensor)
+    result, _ = predict_model(
+        model3d,
+        tensor
+    )
 
     return JSONResponse(result)
 
-
 # =========================================
-# ENSEMBLE
+# ENSEMBLE ENDPOINT
 # =========================================
 @app.post("/predict/ensemble")
-async def predict_ensemble(file: UploadFile = File(...)):
+async def predict_ensemble(
+    file: UploadFile = File(...)
+):
 
-    volume = load_nifti_from_bytes(await file.read())
+    file_bytes = await file.read()
+
+    volume = load_nifti_from_bytes(
+        file_bytes
+    )
 
     tensor2d = preprocess_2d(volume)
+
     tensor3d = preprocess_3d(volume)
 
-    r121, p121 = predict_model(model121, tensor2d)
-    r169, p169 = predict_model(model169, tensor2d)
-    r201, p201 = predict_model(model201, tensor2d)
-    r3d, p3d = predict_model(model3d, tensor3d)
+    r121, p121 = predict_model(
+        model121,
+        tensor2d
+    )
+
+    r169, p169 = predict_model(
+        model169,
+        tensor2d
+    )
+
+    r201, p201 = predict_model(
+        model201,
+        tensor2d
+    )
 
-    avg_probs = (p121 + p169 + p201 + p3d) / 4
+    r3d, p3d = predict_model(
+        model3d,
+        tensor3d
+    )
 
-    pred_idx = torch.argmax(avg_probs).item()
+    avg_probs = (
+        p121 +
+        p169 +
+        p201 +
+        p3d
+    ) / 4
 
-    return {
+    pred_idx = torch.argmax(
+        avg_probs
+    ).item()
 
-        "ensemble_prediction": LABELS[pred_idx],
+    final_result = {
 
-        "ensemble_confidence": round(
-            float(avg_probs[pred_idx]) * 100,
-            2
-        ),
+        "ensemble_prediction":
+            LABELS[pred_idx],
+
+        "ensemble_confidence":
+            round(
+                float(
+                    avg_probs[pred_idx]
+                ) * 100,
+                2
+            ),
 
         "ensemble_probabilities": {
-            LABELS[i]: round(float(avg_probs[i]) * 100, 2)
+
+            LABELS[i]:
+                round(
+                    float(
+                        avg_probs[i]
+                    ) * 100,
+                    2
+                )
+
             for i in range(NUM_CLASSES)
         },
 
         "individual_models": {
 
-            "DenseNet121": r121,
-            "DenseNet169": r169,
-            "DenseNet201": r201,
-            "CNN3D": r3d
+            "DenseNet121":
+                r121,
+
+            "DenseNet169":
+                r169,
+
+            "DenseNet201":
+                r201,
+
+            "CNN3D":
+                r3d
         }
-    }
+    }
+
+    return JSONResponse(final_result)