Hugging Face's logo

Spaces:
Rahma69
/
brainstrokeapp
Runtime error

App Files Community
brainstrokeapp
File size: 4,690 Bytes 
1bfd5f2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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)