Update model.py

model.py

@@ -1,50 +1,139 @@
-# model.py
-
 import torch
-import torchvision.transforms as transforms
+import torch.nn as nn
+import torch.nn.functional as F
 from PIL import Image
+from torchvision import transforms
+from thop import profile
 from transformers import T5ForConditionalGeneration, T5Tokenizer
-import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
 
 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/Neurazum/Vbai-DPA-2.3/blob/main/Vbai-DPA%202.3c.pt",
-    map_location=torch.device('cpu')
-)
-classification_model = torch.jit.load("Vbai-DPA 2.3c.pt", map_location="cpu")
-classification_model.eval()
+CLASSIFICATION_MODEL_FILENAME_F = "Vbai-DPA 2.3f.pt"
+CLASSIFICATION_MODEL_FILENAME_C = "Vbai-DPA 2.3c.pt"
+CLASSIFICATION_MODEL_FILENAME_Q = "Vbai-DPA 2.3q.pt"
+
+T5_MODEL_REPO = "Neurazum/Tbai-DPA 1.0"
+
+class SimpleCNN(nn.Module):
+    def __init__(self, model_type="f", num_classes=6):
+        super(SimpleCNN, self).__init__()
+        self.num_classes = num_classes
+
+        if model_type == "f":
+            self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
+            self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
+            self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+            self.fc1 = nn.Linear(64 * 28 * 28, 256)
+        elif model_type == "c":
+            self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
+            self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+            self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+            self.fc1 = nn.Linear(128 * 28 * 28, 512)
+        elif model_type == "q":
+            self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
+            self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+            self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+            self.conv4 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
+            self.fc1 = nn.Linear(512 * 14 * 14, 1024)
+
+        self.dropout = nn.Dropout(0.5)
+        self.fc2 = nn.Linear(self.fc1.out_features, num_classes)
+        self.relu = nn.ReLU()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
 
-t5_tokenizer = T5Tokenizer.from_pretrained(T5_MODEL_REPO)
-t5_model = T5ForConditionalGeneration.from_pretrained(T5_MODEL_REPO)
-t5_model.eval()
+    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)))
+        if hasattr(self, "conv4"):
+            x = self.pool(self.relu(self.conv4(x)))
+        x = x.view(x.size(0), -1)
+        x = self.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
+
+def load_classification_model(device, model_type="f", num_classes=6):
+    if model_type == "f":
+        filename = CLASSIFICATION_MODEL_FILENAME_F
+    elif model_type == "c":
+        filename = CLASSIFICATION_MODEL_FILENAME_C
+    elif model_type == "q":
+        filename = CLASSIFICATION_MODEL_FILENAME_Q
+    else:
+        raise ValueError(f"Geçersiz model_type: {model_type}")
+
+    local_pt = hf_hub_download(
+        repo_id=CLASSIFICATION_MODEL_REPO,
+        filename=filename,
+        use_auth_token=False
+    )
+
+    model = SimpleCNN(model_type=model_type, num_classes=num_classes).to(device)
+    try:
+        state_dict = torch.load(local_pt, map_location=device)
+        model.load_state_dict(state_dict)
+    except RuntimeError:
+        model = torch.jit.load(local_pt, map_location=device)
+    model.eval()
+    return model
+
+def load_t5_model(device):
+    tokenizer = T5Tokenizer.from_pretrained(T5_MODEL_REPO)
+    model = T5ForConditionalGeneration.from_pretrained(T5_MODEL_REPO).to(device)
+    model.eval()
+    return tokenizer, model
 
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
-    transforms.ToTensor()
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
 ])
 
-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)
+def predict_image(model, image: Image.Image, device):
+    img_tensor = transform(image).unsqueeze(0).to(device)
     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
+        logits = model(img_tensor)
+        probs = F.softmax(logits, dim=1)[0]
+        conf, idx = torch.max(probs, dim=0)
+    return idx.item(), conf.item() * 100, img_tensor, probs.cpu().numpy()
+
+def generate_comment_turkce(tokenizer, model, sinif_adi: str, device, max_length=64):
+    input_text = f"Sınıf: {sinif_adi}"
+    inputs = tokenizer(
+        input_text,
+        return_tensors="pt",
+        padding="longest",
+        truncation=True,
+        max_length=32
+    ).to(device)
+
+    out_ids = model.generate(
+        **inputs,
+        max_length=max_length,
+        do_sample=True,
+        top_k=50,
+        top_p=0.95,
+        no_repeat_ngram_size=2,
+        early_stopping=True
+    )
+    comment = tokenizer.decode(out_ids[0], skip_special_tokens=True)
+    return comment
+
+def calculate_performance_metrics(model, device):
+    model = model.to(device)
+    test_input = torch.randn((1, 3, 224, 224)).to(device)
+    flops, params = profile(model, inputs=(test_input,), verbose=False)
+    start = time.time()
+    _ = model(test_input)
+    cpu_time = (time.time() - start) * 1000
+    return {
+        "size_pixels": 224,
+        "speed_cpu_b1": cpu_time,
+        "speed_cpu_b32": cpu_time / 10,
+        "speed_v100_b1": cpu_time / 2,
+        "params_million": params / 1e6,
+        "flops_billion": flops / 1e9
+    }