initial commit

1 files changed, 37 insertions, 0 deletions

diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..a836ca2
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,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')