Create app.py

app.py

@@ -0,0 +1,150 @@
+# ============================================================
+# BEST FULL PIPELINE – Brain Stroke Classification
+# Dataset: Tekno21 (Normal – Ischemic – Hemorrhagic)
+# Model: EfficientNet (High accuracy)
+# Includes: Training + Accuracy + Error Rate + Gradio UI
+# ============================================================
+
+!pip install -q datasets torch torchvision pillow gradio efficientnet_pytorch
+
+from datasets import load_dataset
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from efficientnet_pytorch import EfficientNet
+from PIL import Image
+import gradio as gr
+import numpy as np
+from sklearn.metrics import accuracy_score
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Device:", device)
+
+# ---------------------- LOAD DATASET ----------------------
+ds = load_dataset("BTX24/tekno21-brain-stroke-dataset-multi")
+
+labels_map = {"NORMAL": 0, "ISCHEMIC": 1, "HEMORRHAGIC": 2}
+
+class_names = ["Normal", "Ischemic", "Hemorrhagic"]
+
+# ---------------------- TRANSFORMS ----------------------
+train_tf = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(10),
+    transforms.ColorJitter(0.2,0.2,0.2,0.1),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+])
+
+val_tf = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+])
+
+# ---------------------- CUSTOM DATASET CLASS ----------------------
+class BrainDataset(Dataset):
+    def __init__(self, hf_data, transform):
+        self.data = hf_data
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        img = Image.open(self.data[idx]["image"]).convert("RGB")
+        label = labels_map[self.data[idx]["label"].upper()]
+        return self.transform(img), label
+
+train_data = BrainDataset(ds["train"], train_tf)
+val_data   = BrainDataset(ds["validation"], val_tf)
+
+train_loader = DataLoader(train_data, batch_size=16, shuffle=True)
+val_loader   = DataLoader(val_data, batch_size=16)
+
+# ---------------------- MODEL ----------------------
+model = EfficientNet.from_pretrained("efficientnet-b0")
+model._fc = nn.Linear(model._fc.in_features, len(class_names))
+model = model.to(device)
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.AdamW(model.parameters(), lr=1e-4)
+
+# ---------------------- TRAINING ----------------------
+EPOCHS = 5
+best_acc = 0
+
+for epoch in range(EPOCHS):
+    model.train()
+    train_correct = 0
+    total = 0
+
+    for imgs, labels in train_loader:
+        imgs, labels = imgs.to(device), labels.to(device)
+        optimizer.zero_grad()
+
+        out = model(imgs)
+        loss = criterion(out, labels)
+
+        loss.backward()
+        optimizer.step()
+
+        _, preds = out.max(1)
+        train_correct += preds.eq(labels).sum().item()
+        total += labels.size(0)
+
+    train_acc = 100 * train_correct / total
+    print(f"Epoch {epoch+1}/{EPOCHS} – Train Acc: {train_acc:.2f}%")
+
+    # Validation
+    model.eval()
+    val_preds = []
+    val_true  = []
+
+    with torch.no_grad():
+        for imgs, labels in val_loader:
+            imgs, labels = imgs.to(device), labels.to(device)
+            out = model(imgs)
+            _, preds = out.max(1)
+
+            val_preds.extend(preds.cpu().numpy())
+            val_true.extend(labels.cpu().numpy())
+
+    val_acc = accuracy_score(val_true, val_preds) * 100
+    print(f"Validation Accuracy: {val_acc:.2f}%")
+
+    if val_acc > best_acc:
+        best_acc = val_acc
+        torch.save(model.state_dict(), "best_model.pth")
+        print("✔ Best Model Saved")
+
+print(" BEST ACCURACY =", best_acc)
+
+# ---------------------- ERROR RATE ----------------------
+error_rate = 100 - best_acc
+print(" ERROR RATE =", error_rate, "%")
+
+# ---------------------- LOAD BEST MODEL ----------------------
+model.load_state_dict(torch.load("best_model.pth"))
+model.eval()
+# ---------------------- GRADIO INTERFACE ----------------------
+def predict(img):
+    img = val_tf(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        out = model(img)
+        probs = torch.softmax(out[0], dim=0).cpu().numpy()
+    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(),
+    title="Brain Stroke Classifier (EfficientNet-B0)",
+    description=f"Best Validation Accuracy: {best_acc:.2f}% | Error Rate: {error_rate:.2f}%"
+)
+
+iface.launch(share=True)