blob: a836ca2e60cf8f64d3a6bdeacc08b03bb21f97d1 (
plain)
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
|
import random
import numpy as np
import torch
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.deterministic = True
torch.backends.cudnn.benchmark = False
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')
|
