Update model.py

model.py

@@ -1,52 +1,50 @@
+# model.py
+
 import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
-
-
-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)
-
-    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)))
-        x = x.view(x.size(0), -1)
-        flattened_size = x.shape[1]
-        self.fc1 = nn.Linear(flattened_size, 512)
-        self.fc2 = nn.Linear(512, num_classes)
-
-    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
-
-def load_model(device: str = 'cpu'):
-    """
-    Downloads and loads the pretrained SimpleCNN model for the 'c' version.
-    """
-    torch_device = torch.device(device)
-
-    weights_path = hf_hub_download(
-        repo_id="Neurazum/Vbai-DPA-2.3",
-        filename="Vbai-DPA 2.3c.pt",
-        repo_type="model"
-    )
-
-    model = SimpleCNN(num_classes=6).to(torch_device)
-    state = torch.load(weights_path, map_location=torch_device)
-    model.load_state_dict(state)
-    model.eval()
-    return model
+import torchvision.transforms as transforms
+from PIL import Image
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+import torch.nn.functional as F
+
+CLASSIFICATION_MODEL_REPO = "Neurazum/Vbai-DPA-2.3"
+T5_MODEL_REPO = "Neurazum/Tbai-DPA-1.0"
+
+classification_model = torch.hub.load_state_dict_from_url(
+    f"https://huggingface.co/{CLASSIFICATION_MODEL_REPO}/resolve/main/vbai_model.pt",
+    map_location=torch.device('cpu')
+)
+classification_model = torch.jit.load("Vbai-DPA 2.3c.pt", map_location="cpu")
+classification_model.eval()
+
+t5_tokenizer = T5Tokenizer.from_pretrained(T5_MODEL_REPO)
+t5_model = T5ForConditionalGeneration.from_pretrained(T5_MODEL_REPO)
+t5_model.eval()
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
+])
+
+class_names = [
+        'Alzheimer Disease',
+        'Mild Alzheimer Risk',
+        'Moderate Alzheimer Risk',
+        'Very Mild Alzheimer Risk',
+        'No Risk',
+        'Parkinson Disease'
+]
+
+def predict(image: Image.Image, question: str = ""):
+    img_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = classification_model(img_tensor)
+        probs = F.softmax(output, dim=1)[0]
+        confidence, pred_idx = torch.max(probs, dim=0)
+        prediction = class_names[pred_idx.item()]
+    
+    input_text = f"Input: {prediction}. Question: {question if question else 'Durum hakkında tıbbi yorum yap'}"
+    t5_input = t5_tokenizer.encode(input_text, return_tensors="pt")
+    t5_output = t5_model.generate(t5_input, max_length=50)
+    comment = t5_tokenizer.decode(t5_output[0], skip_special_tokens=True)
+
+    return prediction, confidence.item(), comment