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
|
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
CLASS_NAMES = [
'Healthy (0)',
'Mild NPDR (1)',
'Moderate NPDR (2)',
'Severe NPDR (3)',
'Proliferative DR (4)',
]
NUM_CLASSES = 5
def set_seed(seed: int = 42) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = True
def compute_class_weights(labels) -> torch.FloatTensor:
"""Inverse-frequency weighting computed from training labels."""
labels = np.array(labels)
class_counts = np.bincount(labels, minlength=NUM_CLASSES)
total = len(labels)
weights = total / (NUM_CLASSES * class_counts.astype(float))
return torch.FloatTensor(weights)
def get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device('cuda')
return torch.device('cpu')
class FocalLoss(nn.Module):
"""Focal loss with optional per-class alpha weighting.
Focuses learning on hard, misclassified examples by down-weighting
easy examples (high pt). Good for class imbalance in medical imaging.
Args:
alpha: Per-class weight tensor (same shape as class weights for CE).
If None, no per-class weighting is applied.
gamma: Focusing parameter. gamma=0 reduces to standard CE.
gamma=2 is the standard value from the RetinaNet paper.
"""
def __init__(self, alpha=None, gamma: float = 2.0):
super().__init__()
self.alpha = alpha
self.gamma = gamma
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
ce_loss = F.cross_entropy(inputs, targets, weight=self.alpha, reduction='none')
pt = torch.exp(-ce_loss)
focal_loss = ((1.0 - pt) ** self.gamma) * ce_loss
return focal_loss.mean()
|
