Tài liệu chi tiết về kiến trúc mô hình AI, cấu hình huấn luyện, dòng chảy dữ liệu (data flow) và hướng dẫn đóng gói suy luận (inference packaging) cho bài toán Phân loại Mức độ Bệnh Võng mạc Tiểu đường (DR - 5 lớp ICDR).
Mô hình dự đoán 5 mức độ tổn thương võng mạc tiểu đường theo tiêu chuẩn quốc tế ICDR:
1flowchart TD
2 A["Input Image (3 x 224 x 224)"] --> B["EfficientNet-B4 Backbone"]
3 B --> C["Feature Maps (1792 x 7 x 7)"]
4
5 subgraph CBAM ["CBAM Attention Module"]
6 C --> D["Channel Attention Module (CA)"]
7 D -->|Feature * CA Map| E["Spatial Attention Module (SA)"]
8 E -->|Feature * SA Map| F["Refined Features (1792 x 7 x 7)"]
9 end
10
11 F --> G["AdaptiveAvgPool2d (1 x 1)"]
12 G --> H["Flatten (1792)"]
13 H --> I["Dropout (p=0.3)"]
14 I --> J["Linear Classification Head (1792 -> 5)"]
15 J --> K["Logits Output (5)"]
16 K --> L["Softmax (Probabilities)"]
-
Channel Attention (Chú ý theo Kênh):
- Gom thông tin không gian bằng
AdaptiveAvgPool2d(1) và AdaptiveMaxPool2d(1).
- Đưa qua shared MLP (2 lớp Conv2d giảm chiều theo tỷ lệ
ratio=16, tức $1792 \rightarrow 112 \rightarrow 1792$).
- Kết hợp kết quả bằng phép cộng, qua hàm
Sigmoid tạo Trọng số Kênh $\mathbf{M}_c$.
- Nhân phần tử với feature map ban đầu.
-
Spatial Attention (Chú ý theo Không gian):
- Gom thông tin kênh bằng phép
mean(dim=1) và max(dim=1) tạo tensor 2 kênh.
- Đưa qua lớp
Conv2d(2 -> 1, kernel_size=7, padding=3) và hàm Sigmoid tạo Trọng số Không gian $\mathbf{M}_s$.
- Nhân phần tử để tập trung vào các vùng tổn thương quan trọng (xuất huyết, vi phình mạch, xuất tiết).
1checkpoint = {
2 "epoch": 24, # Epoch đạt kết quả tốt nhất
3 "model_state_dict": model.state_dict(), # Trọng số tất cả các layer
4 "optimizer_state_dict": optimizer.state_dict(),
5 "val_qwk": 0.8542, # QWK trên tập Validation
6 "val_f1": 0.7215, # Macro F1 trên tập Validation
7 "args": CONFIG # Dictionary lưu tham số cấu hình
8}
Để đưa mô hình này vào ứng dụng sản xuất (Production / Web / Mobile / REST API), bạn cần đóng gói bộ file theo cấu trúc chuẩn bên dưới.
1modelAI_EfficientNetB4/
2├── README.md # Tài liệu cấu hình & kiến trúc (File này)
3├── efficientnet_b4_cbam_fold1.pth # File trọng số PyTorch checkpoint
4├── config.json # Cấu hình nhãn & tiền xử lý
5├── model.py # Đã trích xuất PyTorch class (EfficientNetB4_CBAM)
6├── preprocessing.py # Pipeline tiền xử lý ảnh (Ben Graham + Letterbox)
7├── predictor.py # Class DRPredictor chính để gọi suy luận
8├── gradcam_visualizer.py # Script & module sinh bản đồ nhiệt Grad-CAM
9├── main_api.py # REST API Server với FastAPI (hỗ trợ Grad-CAM)
10├── requirements.txt # Danh sách thư viện phụ thuộc
11├── docs/ # Tài liệu HTML/MD chi tiết
12└── hf_space/ # Gradio UI App cho Hugging Face Space
1import torch
2import torch.nn as nn
3from torchvision import models
4
5class ChannelAttention(nn.Module):
6 def __init__(self, in_planes, ratio=16):
7 super().__init__()
8 self.avg_pool = nn.AdaptiveAvgPool2d(1)
9 self.max_pool = nn.AdaptiveMaxPool2d(1)
10 self.fc = nn.Sequential(
11 nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
12 nn.ReLU(inplace=True),
13 nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False),
14 )
15 self.sigmoid = nn.Sigmoid()
16
17 def forward(self, x):
18 avg_out = self.fc(self.avg_pool(x))
19 max_out = self.fc(self.max_pool(x))
20 return self.sigmoid(avg_out + max_out)
21
22class SpatialAttention(nn.Module):
23 def __init__(self, kernel_size=7):
24 super().__init__()
25 self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
26 self.sigmoid = nn.Sigmoid()
27
28 def forward(self, x):
29 avg_out = torch.mean(x, dim=1, keepdim=True)
30 max_out, _ = torch.max(x, dim=1, keepdim=True)
31 return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))
32
33class CBAM(nn.Module):
34 def __init__(self, in_planes, ratio=16, kernel_size=7):
35 super().__init__()
36 self.ca = ChannelAttention(in_planes, ratio)
37 self.sa = SpatialAttention(kernel_size)
38
39 def forward(self, x):
40 x = x * self.ca(x)
41 x = x * self.sa(x)
42 return x
43
44class EfficientNetB4_CBAM(nn.Module):
45 def __init__(self, num_classes=5, drop_rate=0.3, cbam_ratio=16):
46 super().__init__()
47 backbone = models.efficientnet_b4(weights=None)
48 self.features = backbone.features
49 in_planes = 1792
50 self.cbam = CBAM(in_planes, ratio=cbam_ratio)
51 self.avgpool = nn.AdaptiveAvgPool2d(1)
52 self.classifier = nn.Sequential(
53 nn.Dropout(p=drop_rate),
54 nn.Linear(in_planes, num_classes),
55 )
56
57 def forward(self, x):
58 x = self.features(x)
59 x = self.cbam(x)
60 x = self.avgpool(x)
61 x = torch.flatten(x, 1)
62 return self.classifier(x)
1{
2 "model_name": "EfficientNetB4_CBAM",
3 "num_classes": 5,
4 "input_size": [224, 224],
5 "mean": [0.485, 0.456, 0.406],
6 "std": [0.229, 0.224, 0.225],
7 "labels": {
8 "0": "No DR",
9 "1": "Mild",
10 "2": "Moderate",
11 "3": "Severe",
12 "4": "Proliferative DR"
13 }
14}
1import os
2import json
3import torch
4from PIL import Image
5from torchvision import transforms
6from model import EfficientNetB4_CBAM
7
8class DRPredictor:
9 def __init__(self, weights_path="../efficientnet_b4_cbam_fold1.pth", config_path="config.json"):
10 with open(config_path, "r", encoding="utf-8") as f:
11 self.config = json.load(f)
12
13 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
15 # 1. Khởi tạo kiến trúc
16 self.model = EfficientNetB4_CBAM(
17 num_classes=self.config["num_classes"],
18 drop_rate=0.3
19 )
20
21 # 2. Load trọng số (hỗ trợ cả dạng checkpoint dict hoặc raw state_dict)
22 checkpoint = torch.load(weights_path, map_location=self.device)
23 if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
24 self.model.load_state_dict(checkpoint["model_state_dict"])
25 else:
26 self.model.load_state_dict(checkpoint)
27
28 self.model.to(self.device)
29 self.model.eval()
30
31 # 3. Pipeline Transform chuẩn化
32 self.transform = transforms.Compose([
33 transforms.Resize(tuple(self.config["input_size"])),
34 transforms.ToTensor(),
35 transforms.Normalize(mean=self.config["mean"], std=self.config["std"])
36 ])
37
38 def predict(self, image_path):
39 """
40 Nhận vào đường dẫn ảnh đáy mắt và trả về kết quả phân loại DR.
41 """
42 image = Image.open(image_path).convert("RGB")
43 tensor_img = self.transform(image).unsqueeze(0).to(self.device)
44
45 with torch.no_grad():
46 outputs = self.model(tensor_img)
47 probs = torch.softmax(outputs, dim=1)[0]
48 pred_class = torch.argmax(probs).item()
49
50 return {
51 "class_id": pred_class,
52 "class_name": self.config["labels"][str(pred_class)],
53 "confidence": float(probs[pred_class]),
54 "probabilities": {
55 self.config["labels"][str(i)]: float(probs[i])
56 for i in range(len(probs))
57 }
58 }
1torch>=2.0.0
2torchvision>=0.15.0
3Pillow>=9.5.0
4numpy>=1.24.0
5opencv-python>=4.7.0
