File size: 6,114 Bytes
361b108 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | import numpy as np
import pandas as pd
from torch.utils.data import DataLoader
import torch
import torch.optim as optim
from tqdm import tqdm
import yaml
import wandb
import math
from torch.optim.lr_scheduler import LambdaLR
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.transforms import Activations, AsDiscrete
from monai.networks.nets import SwinUNETR # <-- Ensure you have the correct import for your modelS
from dataloading.dataloader2D import NiftiSegmentationDataset # <-- You must define this properly!
from monai.transforms import (
Activations,
AsDiscrete,
Compose)
import cv2
# -------------------------
# Warmup Cosine Scheduler
# -------------------------
def warmup_cosine_lr_scheduler(optimizer, warmup_epochs, total_epochs):
def lr_lambda(current_epoch):
if current_epoch < warmup_epochs:
return float(current_epoch) / float(max(1, warmup_epochs))
else:
return 0.5 * (1. + math.cos(math.pi * (current_epoch - warmup_epochs) / (total_epochs - warmup_epochs)))
return LambdaLR(optimizer, lr_lambda)
# -------------------------
# Load Config
# -------------------------
with open("/workspace/Segmentation/config2d.yaml", "r") as f:
config = yaml.safe_load(f)
wandb.init(project=config['project_name'], config=config, name=config['run_name'], notes=config['notes'])
cfg = wandb.config
device = torch.device(cfg.device if torch.cuda.is_available() else "cpu")
# -------------------------
# Model Setup
# -------------------------
in_channels = len(cfg.channel_keys) if isinstance(cfg.channel_keys, list) else 1
model =SwinUNETR(img_size=(256, 256), in_channels=in_channels, out_channels=cfg.num_classes, spatial_dims=2).to(device)
optimizer = optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=1e-5)
scheduler = warmup_cosine_lr_scheduler(optimizer, cfg.warmup_epochs, cfg.epochs)
criterion = DiceCELoss(sigmoid=True, to_onehot_y=True)
post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)
# -------------------------
# Dataloaders
# -------------------------
train_dataset = NiftiSegmentationDataset(cfg.csv_file_train, channel_keys=cfg.channel_keys)
train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True)
val_dataset = NiftiSegmentationDataset(cfg.csv_file_val, channel_keys=cfg.channel_keys, augment=False)
val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False)
best_val_loss = float('inf')
# -------------------------
# Training Loop
# -------------------------
for epoch in range(cfg.epochs):
model.train()
total_loss = 0.0
for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{cfg.epochs}"):
x, y = batch
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
logits, _ = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
total_loss += loss.item() * x.size(0)
avg_train_loss = total_loss / len(train_loader.dataset)
# -------------------------
# Validation
# -------------------------
model.eval()
val_loss = 0.0
dice_metric.reset()
with torch.no_grad():
for batch in tqdm(val_loader, desc="Validation"):
x_val, y_val = batch
x_val, y_val = x_val.to(device), y_val.to(device)
val_logits, _ = model(x_val)
v_loss = criterion(val_logits, y_val)
val_loss += v_loss.item() * x_val.size(0)
val_probs = post_pred(val_logits)
dice_metric(y_pred=val_probs, y=y_val)
avg_val_loss = val_loss / len(val_loader.dataset)
avg_dice = dice_metric.aggregate().item()
dice_metric.reset()
# -------------------------
# Logging & Saving
# -------------------------
print(f"Epoch {epoch+1}: Train Loss={avg_train_loss:.4f} | Val Loss={avg_val_loss:.4f} | Dice={avg_dice:.4f}")
if epoch % 30 == 0:
num_examples = min(5, x_val.shape[0])
for i in range(num_examples):
img_np = x_val[i][0].detach().cpu().numpy()
gt_np = y_val[i].detach().cpu().numpy().astype(np.uint8)
pred_np = torch.argmax(val_probs[i], dim=0).detach().cpu().numpy().astype(np.uint8)
# Squeeze image channel if needed
if img_np.shape[0] == 1:
img_np = img_np[0]
# Squeeze and cast to uint8
gt_np = np.squeeze(gt_np).astype(np.uint8)
pred_np = np.squeeze(pred_np).astype(np.uint8)
# Normalize and convert to RGB
img_norm = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
img_uint8 = (img_norm * 255).astype(np.uint8)
img_rgb = cv2.cvtColor(img_uint8, cv2.COLOR_GRAY2BGR)
# Clone image for drawing
contour_img = img_rgb.copy()
# Draw GT contours in green
contours_gt, _ = cv2.findContours(gt_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(contour_img, contours_gt, -1, (0, 255, 0), 2)
# Draw predicted contours in red
contours_pred, _ = cv2.findContours(pred_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(contour_img, contours_pred, -1, (0, 0, 255), 2)
# Log to wandb
wandb.log({
f"val/contours_{i}": wandb.Image(contour_img, caption=f"Epoch {epoch+1} - GT (green), Pred (red)")
})
wandb.log({
"epoch": epoch + 1,
"train_loss": avg_train_loss,
"val_loss": avg_val_loss,
"val_dice": avg_dice,
"lr": scheduler.get_last_lr()[0]
})
scheduler.step()
torch.save(model.state_dict(), '/workspace/Segmentation/checkpoints/latest_model.pth')
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
torch.save(model.state_dict(), '/workspace/Segmentation/checkpoints/best_model.pth')
print("✅ Saved best model.")
|