From 8f571ee51854540d186aafc01d25908bba5ea09e Mon Sep 17 00:00:00 2001 From: Vineet Kumar Date: Sat, 18 Apr 2026 21:34:11 -0400 Subject: initial commit --- .gitignore | 7 +++ channels.scm | 17 ++++++ dataset.py | 123 ++++++++++++++++++++++++++++++++++++++++++ evaluate.py | 146 +++++++++++++++++++++++++++++++++++++++++++++++++ main.py | 161 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ manifest.scm | 13 +++++ models.py | 35 ++++++++++++ requirements.txt | 7 +++ train.py | 136 ++++++++++++++++++++++++++++++++++++++++++++++ utils.py | 37 +++++++++++++ 10 files changed, 682 insertions(+) create mode 100644 .gitignore create mode 100644 channels.scm create mode 100644 dataset.py create mode 100644 evaluate.py create mode 100644 main.py create mode 100644 manifest.scm create mode 100644 models.py create mode 100644 requirements.txt create mode 100644 train.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..69d6571 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +*~ +.claude/ +__pycache__/ +guix-root +guix-root-1-link +results/ +venv/ \ No newline at end of file diff --git a/channels.scm b/channels.scm new file mode 100644 index 0000000..79bf8c0 --- /dev/null +++ b/channels.scm @@ -0,0 +1,17 @@ +(list + (channel + (name 'guix) + (url "https://codeberg.org/dtelsing/Guix") + (branch "pytorch")) + (channel + (name 'guix-rocm) + (url "https://codeberg.org/dtelsing/Guix-ROCm") + (branch "master") + (commit + "84964834d67d3e6edeba5bdef7dc42b4413f91c0")) + (channel + (name 'guix-rocm-ml) + (url "https://codeberg.org/dtelsing/Guix-ROCm-ML") + (branch "master") + (commit + "b6c90eccbbb6e280e7222f12102a69aa963d1de4"))) diff --git a/dataset.py b/dataset.py new file mode 100644 index 0000000..696a312 --- /dev/null +++ b/dataset.py @@ -0,0 +1,123 @@ +import os +from PIL import Image +import numpy as np +import cv2 + +import torch +from torch.utils.data import Dataset, DataLoader +from torchvision import transforms +from sklearn.model_selection import train_test_split + +IMAGENET_MEAN = [0.485, 0.456, 0.406] +IMAGENET_STD = [0.229, 0.224, 0.225] + + +class CLAHETransform: + """Apply CLAHE to the L channel of LAB color space. + + Operates on PIL Images; returns a PIL Image. + """ + + def __init__(self, clip_limit: float = 2.0, tile_grid_size: tuple = (8, 8)): + self.clip_limit = clip_limit + self.tile_grid_size = tile_grid_size + + def __call__(self, img: Image.Image) -> Image.Image: + img_np = np.array(img) # RGB uint8 + lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB) + clahe = cv2.createCLAHE(clipLimit=self.clip_limit, tileGridSize=self.tile_grid_size) + lab[:, :, 0] = clahe.apply(lab[:, :, 0]) + result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB) + return Image.fromarray(result) + + +class DRDataset(Dataset): + """Diabetic Retinopathy dataset from flat class-folder layout.""" + + def __init__(self, image_paths: list, labels: list, transform=None): + self.image_paths = image_paths + self.labels = labels + self.transform = transform + + def __len__(self) -> int: + return len(self.image_paths) + + def __getitem__(self, idx: int): + img = Image.open(self.image_paths[idx]).convert('RGB') + label = self.labels[idx] + if self.transform is not None: + img = self.transform(img) + return img, label + + +def split_dataset(data_root: str, seed: int = 42): + """Scan class folders 0–4 and return stratified 70/15/15 splits. + + Returns: + (train_paths, train_labels, val_paths, val_labels, test_paths, test_labels) + """ + all_paths, all_labels = [], [] + for class_idx in range(5): + class_dir = os.path.join(data_root, str(class_idx)) + for fname in sorted(os.listdir(class_dir)): + if fname.lower().endswith(('.jpg', '.jpeg', '.png')): + all_paths.append(os.path.join(class_dir, fname)) + all_labels.append(class_idx) + + # 70% train, 30% temp + train_paths, temp_paths, train_labels, temp_labels = train_test_split( + all_paths, all_labels, test_size=0.30, stratify=all_labels, random_state=seed + ) + # 50/50 of temp -> 15% val, 15% test + val_paths, test_paths, val_labels, test_labels = train_test_split( + temp_paths, temp_labels, test_size=0.50, stratify=temp_labels, random_state=seed + ) + + print(f"Dataset split — train: {len(train_paths)}, val: {len(val_paths)}, test: {len(test_paths)}") + return train_paths, train_labels, val_paths, val_labels, test_paths, test_labels + + +def build_transforms(use_clahe: bool, training: bool): + """Build a transform pipeline for training or eval.""" + ops = [transforms.Resize((224, 224))] + if use_clahe: + ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8))) + if training: + ops += [ + transforms.RandomHorizontalFlip(p=0.5), + transforms.RandomVerticalFlip(p=0.5), + transforms.RandomRotation(degrees=15), + ] + ops += [ + transforms.ToTensor(), + transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), + ] + return transforms.Compose(ops) + + +def build_dataloaders( + train_paths, train_labels, + val_paths, val_labels, + test_paths, test_labels, + use_clahe: bool, + batch_size: int = 32, + num_workers: int = 4, +): + """Create DataLoaders for all three splits.""" + train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True)) + val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False)) + test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False)) + + train_loader = DataLoader( + train_ds, batch_size=batch_size, shuffle=True, + num_workers=num_workers, pin_memory=True + ) + val_loader = DataLoader( + val_ds, batch_size=batch_size * 2, shuffle=False, + num_workers=num_workers, pin_memory=True + ) + test_loader = DataLoader( + test_ds, batch_size=batch_size * 2, shuffle=False, + num_workers=num_workers, pin_memory=True + ) + return train_loader, val_loader, test_loader diff --git a/evaluate.py b/evaluate.py new file mode 100644 index 0000000..7a39db4 --- /dev/null +++ b/evaluate.py @@ -0,0 +1,146 @@ +import os +import csv + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +import numpy as np +import matplotlib.pyplot as plt +import seaborn as sns +from sklearn.metrics import ( + accuracy_score, + f1_score, + classification_report, + confusion_matrix, + cohen_kappa_score, +) + +from utils import CLASS_NAMES + + +def evaluate_model(model: nn.Module, test_loader: DataLoader, device: torch.device) -> dict: + """Run inference on test_loader and return a dict of all metrics.""" + model.eval() + all_preds, all_labels = [], [] + + with torch.no_grad(): + for images, labels in test_loader: + images = images.to(device, non_blocking=True) + outputs = model(images) + preds = outputs.argmax(dim=1) + all_preds.extend(preds.cpu().numpy()) + all_labels.extend(labels.numpy()) + + all_preds = np.array(all_preds) + all_labels = np.array(all_labels) + + return { + 'accuracy': accuracy_score(all_labels, all_preds), + 'weighted_f1': f1_score(all_labels, all_preds, average='weighted', zero_division=0), + 'macro_f1': f1_score(all_labels, all_preds, average='macro', zero_division=0), + 'classification_report': classification_report( + all_labels, all_preds, target_names=CLASS_NAMES, zero_division=0 + ), + 'confusion_matrix': confusion_matrix(all_labels, all_preds), + 'cohen_kappa': cohen_kappa_score(all_labels, all_preds, weights='quadratic'), + } + + +def plot_confusion_matrix(cm: np.ndarray, save_path: str) -> None: + """Save a seaborn confusion matrix heatmap to save_path.""" + fig, ax = plt.subplots(figsize=(8, 6)) + sns.heatmap( + cm, annot=True, fmt='d', cmap='Blues', + xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES, + ax=ax + ) + ax.set_xlabel('Predicted') + ax.set_ylabel('True') + ax.set_title('Confusion Matrix') + plt.tight_layout() + fig.savefig(save_path, dpi=150) + plt.close(fig) + + +def plot_training_curves( + train_losses: list, val_losses: list, val_f1s: list, save_path: str +) -> None: + """Save loss and F1 training curves to save_path.""" + epochs = range(1, len(train_losses) + 1) + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) + + ax1.plot(epochs, train_losses, label='Train loss') + ax1.plot(epochs, val_losses, label='Val loss') + ax1.set_xlabel('Epoch') + ax1.set_ylabel('Loss') + ax1.set_title('Training and Validation Loss') + ax1.legend() + + ax2.plot(epochs, val_f1s, label='Val weighted F1', color='green') + ax2.set_xlabel('Epoch') + ax2.set_ylabel('Weighted F1') + ax2.set_title('Validation Weighted F1') + ax2.legend() + + plt.tight_layout() + fig.savefig(save_path, dpi=150) + plt.close(fig) + + +def generate_summary(all_results: dict, save_dir: str) -> None: + """Print comparison table, save as CSV, and generate bar chart.""" + os.makedirs(save_dir, exist_ok=True) + + rows = [] + for key, r in all_results.items(): + model_name, clahe_str = key.rsplit('_clahe=', 1) + rows.append({ + 'model': model_name, + 'clahe': clahe_str, + 'weighted_f1': r['weighted_f1'], + 'macro_f1': r['macro_f1'], + 'accuracy': r['accuracy'], + 'cohen_kappa': r['cohen_kappa'], + 'train_time_s': round(r['train_time'], 1), + 'epochs_trained': r['epochs_trained'], + }) + + # Print table + header = f"{'Model':<18} {'CLAHE':<6} {'W-F1':>6} {'M-F1':>6} {'Acc':>6} {'Kappa':>6} {'Time(s)':>8} {'Epochs':>7}" + print('\n' + '=' * len(header)) + print(header) + print('-' * len(header)) + for row in rows: + print( + f"{row['model']:<18} {row['clahe']:<6} " + f"{row['weighted_f1']:>6.4f} {row['macro_f1']:>6.4f} " + f"{row['accuracy']:>6.4f} {row['cohen_kappa']:>6.4f} " + f"{row['train_time_s']:>8} {row['epochs_trained']:>7}" + ) + print('=' * len(header)) + + # Save CSV + csv_path = os.path.join(save_dir, 'summary.csv') + with open(csv_path, 'w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + print(f"\nSummary saved to {csv_path}") + + # Comparison bar chart + labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows] + w_f1_vals = [r['weighted_f1'] for r in rows] + + fig, ax = plt.subplots(figsize=(10, 5)) + bars = ax.bar(labels, w_f1_vals, color=['#4C72B0', '#DD8452'] * 3) + ax.set_ylabel('Weighted F1-Score') + ax.set_title('Weighted F1 Comparison Across Experiments') + ax.set_ylim(0, 1.0) + for bar, val in zip(bars, w_f1_vals): + ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01, + f'{val:.4f}', ha='center', va='bottom', fontsize=9) + plt.tight_layout() + chart_path = os.path.join(save_dir, 'comparison_chart.png') + fig.savefig(chart_path, dpi=150) + plt.close(fig) + print(f"Comparison chart saved to {chart_path}") diff --git a/main.py b/main.py new file mode 100644 index 0000000..6d8b2d9 --- /dev/null +++ b/main.py @@ -0,0 +1,161 @@ +""" +EEL4759 Final Project — Diabetic Retinopathy Classification +Compares ResNet-50, EfficientNet-B0, and ViT-B/16 with and without CLAHE preprocessing. + +Usage: + python main.py # run all 6 experiments (default) + python main.py --models resnet50 --clahe 0 --epochs 2 # quick sanity check +""" + +import argparse +import os + +import torch +import torch.nn as nn +from torch.optim.lr_scheduler import CosineAnnealingLR +import kagglehub + +from utils import set_seed, compute_class_weights, get_device, CLASS_NAMES +from dataset import split_dataset, build_dataloaders +from models import create_model +from train import train_model +from evaluate import evaluate_model, plot_confusion_matrix, plot_training_curves, generate_summary + +RESULTS_DIR = 'results' + + +def parse_args(): + parser = argparse.ArgumentParser(description='Diabetic Retinopathy Classification') + parser.add_argument( + '--models', nargs='+', + default=['resnet50', 'efficientnet_b0', 'vit_b_16'], + help='Models to run' + ) + parser.add_argument( + '--clahe', nargs='+', type=int, default=[0, 1], + help='CLAHE settings to run (0=off, 1=on)' + ) + parser.add_argument('--epochs', type=int, default=30, help='Max training epochs') + parser.add_argument('--patience', type=int, default=5, help='Early stopping patience') + parser.add_argument('--batch-size', type=int, default=32, help='Training batch size') + parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') + parser.add_argument('--weight-decay', type=float, default=1e-4, help='AdamW weight decay') + parser.add_argument('--seed', type=int, default=42, help='Random seed') + parser.add_argument('--data-root', type=str, default=None, + help='Path to dataset root (auto-downloaded if not set)') + parser.add_argument('--num-workers', type=int, default=4, help='DataLoader workers') + return parser.parse_args() + + +def main(): + args = parse_args() + set_seed(args.seed) + device = get_device() + print(f"Using device: {device}") + + # ── Dataset ────────────────────────────────────────────────────────────── + if args.data_root is None: + print("Downloading/locating dataset via kagglehub...") + data_root = kagglehub.dataset_download("amanneo/diabetic-retinopathy-resized-arranged") + else: + data_root = args.data_root + print(f"Dataset root: {data_root}") + + train_paths, train_labels, val_paths, val_labels, test_paths, test_labels = \ + split_dataset(data_root, seed=args.seed) + + # Class weights computed once from training labels + class_weights = compute_class_weights(train_labels).to(device) + print("Class weights:", [f"{w:.3f}" for w in class_weights.cpu().numpy()]) + + # ── Experiments ────────────────────────────────────────────────────────── + experiments = [ + (model_name, bool(use_clahe)) + for model_name in args.models + for use_clahe in args.clahe + ] + + all_results = {} + os.makedirs(RESULTS_DIR, exist_ok=True) + + for model_name, use_clahe in experiments: + exp_key = f"{model_name}_clahe={use_clahe}" + exp_dir = os.path.join(RESULTS_DIR, exp_key) + os.makedirs(exp_dir, exist_ok=True) + + print(f"\n{'='*70}") + print(f"Experiment: {model_name} | CLAHE: {use_clahe}") + print(f"{'='*70}") + + train_loader, val_loader, test_loader = build_dataloaders( + train_paths, train_labels, + val_paths, val_labels, + test_paths, test_labels, + use_clahe=use_clahe, + batch_size=args.batch_size, + num_workers=args.num_workers, + ) + + model = create_model(model_name, num_classes=5, pretrained=True).to(device) + criterion = nn.CrossEntropyLoss(weight=class_weights) + optimizer = torch.optim.AdamW( + model.parameters(), lr=args.lr, weight_decay=args.weight_decay + ) + scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs) + + # Train + train_results = train_model( + model, train_loader, val_loader, + criterion, optimizer, scheduler, + device, + num_epochs=args.epochs, + patience=args.patience, + model_name=exp_key, + ) + + # Save best checkpoint + ckpt_path = os.path.join(exp_dir, 'best_model.pth') + torch.save(train_results['best_model_state'], ckpt_path) + print(f" Checkpoint saved: {ckpt_path}") + + # Load best weights for evaluation + model.load_state_dict(train_results['best_model_state']) + + # Evaluate on test set + test_metrics = evaluate_model(model, test_loader, device) + + print(f"\n Test results — Weighted F1: {test_metrics['weighted_f1']:.4f} | " + f"Macro F1: {test_metrics['macro_f1']:.4f} | " + f"Accuracy: {test_metrics['accuracy']:.4f} | " + f"Kappa: {test_metrics['cohen_kappa']:.4f}") + print(test_metrics['classification_report']) + + # Save per-class report + report_path = os.path.join(exp_dir, 'classification_report.txt') + with open(report_path, 'w') as f: + f.write(f"Experiment: {exp_key}\n\n") + f.write(test_metrics['classification_report']) + + # Plots + plot_confusion_matrix( + test_metrics['confusion_matrix'], + save_path=os.path.join(exp_dir, 'confusion_matrix.png'), + ) + plot_training_curves( + train_results['train_losses'], + train_results['val_losses'], + train_results['val_f1s'], + save_path=os.path.join(exp_dir, 'training_curves.png'), + ) + + all_results[exp_key] = {**train_results, **test_metrics} + # Remove bulky state dict from summary dict (already saved to disk) + all_results[exp_key].pop('best_model_state', None) + + # ── Summary ─────────────────────────────────────────────────────────────── + if len(all_results) > 1: + generate_summary(all_results, save_dir=RESULTS_DIR) + + +if __name__ == '__main__': + main() diff --git a/manifest.scm b/manifest.scm new file mode 100644 index 0000000..38c3cdf --- /dev/null +++ b/manifest.scm @@ -0,0 +1,13 @@ +(module-define! + (resolve-module '(guix-rocm packages rocm-base)) + '%amdgpu-targets + (lambda () + "gfx1100")) + +(specifications->manifest + (list "python" + "python-flask" + "python-pytorch-rocm" + "python-torchvision-rocm" + "rocminfo" + "rocm-smi")) diff --git a/models.py b/models.py new file mode 100644 index 0000000..ab21055 --- /dev/null +++ b/models.py @@ -0,0 +1,35 @@ +import torch.nn as nn +from torchvision.models import ( + resnet50, ResNet50_Weights, + efficientnet_b0, EfficientNet_B0_Weights, + vit_b_16, ViT_B_16_Weights, +) + + +def create_model(model_name: str, num_classes: int = 5, pretrained: bool = True) -> nn.Module: + """Create a pretrained model with the classification head replaced for num_classes. + + Supported model_name values: 'resnet50', 'efficientnet_b0', 'vit_b_16' + All backbone layers are trainable (full fine-tuning). + """ + if model_name == 'resnet50': + weights = ResNet50_Weights.DEFAULT if pretrained else None + model = resnet50(weights=weights) + model.fc = nn.Linear(model.fc.in_features, num_classes) # 2048 -> num_classes + + elif model_name == 'efficientnet_b0': + weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None + model = efficientnet_b0(weights=weights) + # classifier is Sequential(Dropout(0.2), Linear(1280, 1000)) + model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes) + + elif model_name == 'vit_b_16': + weights = ViT_B_16_Weights.DEFAULT if pretrained else None + model = vit_b_16(weights=weights) + # heads is Sequential(head=Linear(768, 1000)) + model.heads.head = nn.Linear(model.heads.head.in_features, num_classes) + + else: + raise ValueError(f"Unknown model: {model_name!r}. Choose from: resnet50, efficientnet_b0, vit_b_16") + + return model diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e77f7b5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +kagglehub +torch +torchvision +scikit-learn +opencv-python-headless +matplotlib +seaborn diff --git a/train.py b/train.py new file mode 100644 index 0000000..545b299 --- /dev/null +++ b/train.py @@ -0,0 +1,136 @@ +import copy +import time + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from sklearn.metrics import f1_score + + +def train_model( + model: nn.Module, + train_loader: DataLoader, + val_loader: DataLoader, + criterion: nn.Module, + optimizer: torch.optim.Optimizer, + scheduler, + device: torch.device, + num_epochs: int = 30, + patience: int = 5, + model_name: str = "", +) -> dict: + """Train model with mixed precision, early stopping on val weighted-F1. + + Returns a dict with: + best_model_state, train_losses, val_losses, val_f1s, + train_time (seconds), epochs_trained + """ + scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None + + best_val_f1 = -1.0 + best_state = None + patience_counter = 0 + + train_losses, val_losses, val_f1s = [], [], [] + start_time = time.time() + + for epoch in range(1, num_epochs + 1): + # ── Training phase ────────────────────────────────────────────────── + model.train() + running_loss = 0.0 + n_train = 0 + + for images, labels in train_loader: + images = images.to(device, non_blocking=True) + labels = labels.to(device, non_blocking=True) + + optimizer.zero_grad() + + if scaler is not None: + with torch.amp.autocast('cuda'): + outputs = model(images) + loss = criterion(outputs, labels) + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + scaler.step(optimizer) + scaler.update() + else: + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + optimizer.step() + + running_loss += loss.item() * images.size(0) + n_train += images.size(0) + + train_loss = running_loss / n_train + train_losses.append(train_loss) + + # ── Validation phase ───────────────────────────────────────────────── + model.eval() + val_running_loss = 0.0 + n_val = 0 + all_preds, all_labels = [], [] + + with torch.no_grad(): + for images, labels in val_loader: + images = images.to(device, non_blocking=True) + labels = labels.to(device, non_blocking=True) + + if scaler is not None: + with torch.amp.autocast('cuda'): + outputs = model(images) + loss = criterion(outputs, labels) + else: + outputs = model(images) + loss = criterion(outputs, labels) + + val_running_loss += loss.item() * images.size(0) + n_val += images.size(0) + + preds = outputs.argmax(dim=1) + all_preds.extend(preds.cpu().numpy()) + all_labels.extend(labels.cpu().numpy()) + + val_loss = val_running_loss / n_val + val_losses.append(val_loss) + + val_f1 = f1_score(all_labels, all_preds, average='weighted', zero_division=0) + val_f1s.append(val_f1) + + current_lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else optimizer.param_groups[0]['lr'] + print( + f"[{model_name}] Epoch {epoch:3d}/{num_epochs} | " + f"train loss: {train_loss:.4f} | val loss: {val_loss:.4f} | " + f"val F1: {val_f1:.4f} | lr: {current_lr:.2e}" + ) + + scheduler.step() + + # ── Early stopping ──────────────────────────────────────────────────── + if val_f1 > best_val_f1: + best_val_f1 = val_f1 + best_state = copy.deepcopy(model.state_dict()) + patience_counter = 0 + else: + patience_counter += 1 + if patience_counter >= patience: + print(f" Early stopping at epoch {epoch} (no improvement for {patience} epochs).") + break + + train_time = time.time() - start_time + epochs_trained = len(train_losses) + + print(f" Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val F1: {best_val_f1:.4f}") + + return { + 'best_model_state': best_state, + 'train_losses': train_losses, + 'val_losses': val_losses, + 'val_f1s': val_f1s, + 'train_time': train_time, + 'epochs_trained': epochs_trained, + 'best_val_f1': best_val_f1, + } 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') -- cgit v1.2.3