utils.py (1892B)
1 import random 2 import numpy as np 3 import torch 4 import torch.nn as nn 5 import torch.nn.functional as F 6 7 CLASS_NAMES = [ 8 'Healthy (0)', 9 'Mild NPDR (1)', 10 'Moderate NPDR (2)', 11 'Severe NPDR (3)', 12 'Proliferative DR (4)', 13 ] 14 15 NUM_CLASSES = 5 16 17 18 def set_seed(seed: int = 42) -> None: 19 random.seed(seed) 20 np.random.seed(seed) 21 torch.manual_seed(seed) 22 torch.cuda.manual_seed_all(seed) 23 torch.backends.cudnn.benchmark = True 24 25 26 def compute_class_weights(labels) -> torch.FloatTensor: 27 """Inverse-frequency weighting computed from training labels.""" 28 labels = np.array(labels) 29 class_counts = np.bincount(labels, minlength=NUM_CLASSES) 30 total = len(labels) 31 weights = total / (NUM_CLASSES * class_counts.astype(float)) 32 return torch.FloatTensor(weights) 33 34 35 def get_device() -> torch.device: 36 if torch.cuda.is_available(): 37 return torch.device('cuda') 38 return torch.device('cpu') 39 40 41 class FocalLoss(nn.Module): 42 """Focal loss with optional per-class alpha weighting. 43 44 Focuses learning on hard, misclassified examples by down-weighting 45 easy examples (high pt). Good for class imbalance in medical imaging. 46 47 Args: 48 alpha: Per-class weight tensor (same shape as class weights for CE). 49 If None, no per-class weighting is applied. 50 gamma: Focusing parameter. gamma=0 reduces to standard CE. 51 gamma=2 is the standard value from the RetinaNet paper. 52 """ 53 54 def __init__(self, alpha=None, gamma: float = 2.0): 55 super().__init__() 56 self.alpha = alpha 57 self.gamma = gamma 58 59 def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: 60 ce_loss = F.cross_entropy(inputs, targets, weight=self.alpha, reduction='none') 61 pt = torch.exp(-ce_loss) 62 focal_loss = ((1.0 - pt) ** self.gamma) * ce_loss 63 return focal_loss.mean()