Update app.py

app.py

@@ -1,45 +1,43 @@
+import gradio as gr
 import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
+from torchvision import transforms
+from PIL import Image
+from model import load_model
+import numpy as np
 
-class SimpleCNN(nn.Module):
-    def __init__(self, num_classes=6):
-        super(SimpleCNN, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
-        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
-        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.relu = nn.ReLU()
-        self.dropout = nn.Dropout(0.5)
-        self._initialize_fc(num_classes)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = load_model(device)
 
-    def _initialize_fc(self, num_classes):
-        dummy_input = torch.zeros(1, 3, 448, 448)
-        x = self.pool(self.relu(self.conv1(dummy_input)))
-        x = self.pool(self.relu(self.conv2(x)))
-        x = self.pool(self.relu(self.conv3(x)))
-        flattened_size = x.view(x.size(0), -1).shape[1]
-        self.fc1 = nn.Linear(flattened_size, 512)
-        self.fc2 = nn.Linear(512, num_classes)
+class_names = [
+    'Alzheimer Disease', 
+    'Mild Alzheimer Risk', 
+    'Moderate Alzheimer Risk',
+    'Very Mild Alzheimer Risk', 
+    'No Risk', 
+    'Parkinson Disease'
+]
 
-    def forward(self, x):
-        x = self.pool(self.relu(self.conv1(x)))
-        x = self.pool(self.relu(self.conv2(x)))
-        x = self.pool(self.relu(self.conv3(x)))
-        x = x.view(x.size(0), -1)
-        x = self.dropout(self.relu(self.fc1(x)))
-        x = self.fc2(x)
-        return x
+transform = transforms.Compose([
+    transforms.Resize((448, 448)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
 
+def predict(image):
+    image = image.convert("RGB")
+    tensor = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = model(tensor)
+        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
+        predicted = torch.argmax(probs).item()
+        confidence = probs[predicted].item() * 100
+    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
 
-def load_model(device: str = 'cpu'):
-    weights_path = hf_hub_download(
-        repo_id="Neurazum/Vbai-DPA-2.2",
-        filename="Vbai-DPA 2.2c.pt",
-        repo_type="model"
-    )
-    model = SimpleCNN(num_classes=6).to(device)
-    state = torch.load(weights_path, map_location=device)
-    model.load_state_dict(state)
-    model.eval()
-    return model
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Vbai-DPA 2.2 (C Version)",
+    description="Upload an MRI and fMRI image to classify the risk level using the 'C' version of the Vbai-DPA 2.2 model."
+).launch()