train.py

import os
import cv2
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from tqdm import tqdm

from unet.unet_model import UNet


# =========================
# Dataset
# =========================
class SegDataset(Dataset):
    def __init__(self, img_dir, mask_dir):
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.files = sorted(os.listdir(img_dir))

        self.transform = transforms.Compose([
            transforms.ToTensor()
        ])

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        fname = self.files[idx]

        img = cv2.imread(os.path.join(self.img_dir, fname), 0)

        mask = cv2.imread(os.path.join(self.mask_dir, fname), 0)
        
        img = torch.from_numpy(img).float().unsqueeze(0) / 255.0
        mask = (mask > 127).astype(np.float32)
        mask = torch.from_numpy(mask).unsqueeze(0)

        return img, mask


# =========================
# Training setup
# =========================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Training on {device}")

dataset = SegDataset(
    "data/train_img",
    "data/train_lab"
)

val_dataset = SegDataset(
    "data/test_img",
    "data/test_lab"
)

loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=0)

model = UNet(n_channels=1, n_classes=1)
model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = torch.nn.BCEWithLogitsLoss()

epochs = 30


# =========================
# Training loop
# =========================
for epoch in range(epochs):

    model.train()
    epoch_loss = 0

    for imgs, masks in tqdm(loader):

        imgs = imgs.to(device)
        masks = masks.to(device)

        preds = model(imgs)

        loss = criterion(preds, masks)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    # Validation
    model.eval()
    val_loss = 0
    
    with torch.no_grad():
        for imgs, masks in tqdm(val_loader, desc="Validation"):
            imgs = imgs.to(device)
            masks = masks.to(device)
            
            preds = model(imgs)
            loss = criterion(preds, masks)
            val_loss += loss.item()

    print(f"Epoch {epoch} Loss: {epoch_loss/len(loader):.4f} | Val Loss: {val_loss/len(val_loader):.4f}")

    torch.save(model.state_dict(), "unet_model.pth")