""" SegFormer Fine-tuning Script This script fine-tunes a SegFormer model on a custom semantic segmentation dataset. It provides configurable parameters for training hyperparameters and dataset settings. """ import json import os import zipfile import numpy as np import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from PIL import Image from pathlib import Path from datetime import datetime from transformers import ( SegformerImageProcessor, SegformerForSemanticSegmentation, ) import evaluate from tqdm import tqdm class SemanticSegmentationDataset(Dataset): """Image (semantic) segmentation dataset.""" def __init__( self, root_dir, image_processor, train=True, data_percent=100, ): """ Args: root_dir (string): Root directory of the dataset containing the images + annotations. image_processor (SegFormerImageProcessor): image processor to prepare images + segmentation maps. train (bool): Whether to load "training" or "validation" images + annotations. data_percent (int): Percentage of the dataset to use. 100 means all data, 50 means half of the data. """ self.root_dir = root_dir self.image_processor = image_processor self.train = train sub_path = "training" if self.train else "validation" self.img_dir = os.path.join(self.root_dir, "images", sub_path) self.ann_dir = os.path.join(self.root_dir, "annotations", sub_path) # read images image_file_names = [] for root, dirs, files in os.walk(self.img_dir): image_file_names.extend(files) self.images = sorted(image_file_names) # read annotations annotation_file_names = [] for root, dirs, files in os.walk(self.ann_dir): annotation_file_names.extend(files) self.annotations = sorted(annotation_file_names) assert len(self.images) == len( self.annotations ), "There must be as many images as there are segmentation maps" # Apply data_percent to limit the dataset size data_percent = data_percent / 100.0 if data_percent < 1.0: images_num_samples = int(len(self.images) * data_percent) annotations_num_samples = int(len(self.annotations) * data_percent) self.images = self.images[:images_num_samples] self.annotations = self.annotations[:annotations_num_samples] def __len__(self): return len(self.images) def __getitem__(self, idx): image = Image.open(os.path.join(self.img_dir, self.images[idx])) segmentation_map = Image.open( os.path.join( self.ann_dir, self.annotations[idx], ), ) encoded_inputs = self.image_processor( image, segmentation_map, return_tensors="pt", ) for k, v in encoded_inputs.items(): encoded_inputs[k].squeeze_() # remove batch dimension return encoded_inputs class MeanDice: def __init__(self): self.reset() def reset(self): """Reset stored predictions and references.""" self.predictions = [] self.references = [] def add_batch(self, predictions, references): """ Add a batch of predictions and references. Args: predictions (np.ndarray): Predicted class indices references (np.ndarray): Ground truth class indices """ self.predictions.append(predictions) self.references.append(references) def compute(self, num_labels, ignore_index=None): """Compute mean Dice score across all stored batches.""" predictions = np.concatenate([p.flatten() for p in self.predictions]) references = np.concatenate([r.flatten() for r in self.references]) dice_scores = [] for class_id in range(num_labels): pred_mask = predictions == class_id ref_mask = references == class_id # Exclude ignore_index if ignore_index is not None: valid_mask = references != ignore_index pred_mask = pred_mask & valid_mask ref_mask = ref_mask & valid_mask intersection = np.sum(pred_mask & ref_mask) union = np.sum(pred_mask) + np.sum(ref_mask) if union == 0: dice = 1.0 if intersection == 0 else 0.0 else: dice = 2.0 * intersection / union dice_scores.append(dice) return { "mean_dice": float(np.mean(dice_scores)), "per_class_dice": dice_scores, } def get_latest_model_dir(base_path: str = "./segformer_finetuned") -> Path: """ Returns the Path to the latest model directory based on timestamp folder names. Folder names must follow the format: YYYY-MM-DD_HH-MM-SS """ base = Path(base_path) if not base.exists() or not base.is_dir(): raise FileNotFoundError(f"Directory not found: {base_path}") model_dirs = [] for d in base.iterdir(): if d.is_dir(): try: dt = datetime.strptime(d.name, "%Y-%m-%d_%H-%M-%S") model_dirs.append((dt, d)) except ValueError: continue # Skip non-matching directories if not model_dirs: raise FileNotFoundError( "No model directories found with valid timestamp format." ) # Return the directory with the latest timestamp return max(model_dirs, key=lambda x: x[0])[1] def load_model_and_labels(data_dir, model_path): """Load the model and label mappings.""" # Load id2label mapping from JSON file id2label = json.load(open(f"{data_dir}/id2label.json", mode="r")) id2label = {int(k): v for k, v in id2label.items()} label2id = {v: k for k, v in id2label.items()} # Load id2color mapping from JSON file id2color = json.load(open(f"{data_dir}/id2color.json", "r")) print(f"Loaded {len(id2label)} classes:") for i, label in id2label.items(): print(f" {i}: {label}") # Load model model = SegformerForSemanticSegmentation.from_pretrained( model_path, num_labels=len(id2label), id2label=id2label, label2id=label2id, ) return model, id2label, id2color def create_datasets_and_dataloaders( image_width, image_height, data_dir, batch_size, data_percent, ): """Create datasets and dataloaders.""" image_processor = SegformerImageProcessor( size={"height": image_height, "width": image_width}, ) train_dataset = SemanticSegmentationDataset( root_dir=data_dir, image_processor=image_processor, train=True, data_percent=data_percent, ) valid_dataset = SemanticSegmentationDataset( root_dir=data_dir, image_processor=image_processor, train=False, data_percent=data_percent, ) print(f"Number of training examples: {len(train_dataset)}") print(f"Number of validation examples: {len(valid_dataset)}") train_dataloader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, ) valid_dataloader = DataLoader( valid_dataset, batch_size=batch_size, ) return train_dataloader, valid_dataloader def class_indices_to_rgb(class_indices, id2color): """Convert class indices to RGB colored image.""" # class_indices shape: (H, W) with integer class IDs height, width = class_indices.shape rgb_image = np.zeros((height, width, 3), dtype=np.uint8) for class_id, color in id2color.items(): rgb_image[class_indices == class_id] = color return rgb_image def validate_model( model: SegformerForSemanticSegmentation, dataloader, device, id2label, calc_dice=False, epoch=None, ): """ Validate the model on a validation set and return loss, IoU, accuracy. """ model.eval() metric = evaluate.load("mean_iou") dice = MeanDice() total_loss = 0.0 num_batches = 0 with torch.no_grad(): for batch in tqdm( dataloader, desc="Validating Epoch " + str(epoch if epoch is not None else ""), leave=False, unit="batches", ): pixel_values = batch["pixel_values"].to(device) labels = batch["labels"].to(device) outputs = model(pixel_values=pixel_values, labels=labels) logits = outputs.logits loss = outputs.loss total_loss += loss.item() num_batches += 1 upsampled_logits = nn.functional.interpolate( logits, size=labels.shape[-2:], mode="bilinear", align_corners=False, ) predicted = upsampled_logits.argmax(dim=1) # Store predictions and references for additional metrics pred_np = predicted.detach().cpu().numpy() ref_np = labels.detach().cpu().numpy() metric.add_batch( predictions=pred_np, references=ref_np, ) if calc_dice: dice.add_batch( predictions=pred_np, references=ref_np, ) # Calculate IoU and accuracy result = metric.compute( num_labels=len(id2label), ignore_index=10, reduce_labels=False, ) if calc_dice: dice_result = dice.compute( num_labels=len(id2label), ignore_index=10, ) avg_loss = total_loss / num_batches if num_batches > 0 else 0.0 return ( avg_loss, result["mean_iou"], result["per_category_iou"], result["mean_accuracy"], result["per_category_accuracy"], dice_result["mean_dice"] if calc_dice else None, dice_result["per_class_dice"] if calc_dice else None, ) def run_training( model: SegformerForSemanticSegmentation, device, train_dataloader, valid_dataloader, id2label, num_epochs, learning_rate, early_stopping, validate_every, ): """Train the model. Returns ------- tuple(best_model, metrics) best_model : nn.Module metrics : dict with lists for keys: 'epoch', 'train_loss', 'train_iou', 'train_acc', 'val_loss', 'val_iou', 'val_acc' """ # Setup device model.to(device) # Setup optimizer optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) # Setup metrics metrics = { "epoch": [], "train_loss": [], "train_iou": [], "train_acc": [], "val_loss": [], "val_iou": [], "val_acc": [], } metric = evaluate.load("mean_iou") model.train() # Initial validation ( loss, iou, per_class_iou, acc, per_class_acc, dice, dice_per_class, ) = validate_model( model=model, dataloader=valid_dataloader, device=device, id2label=id2label, calc_dice=True, epoch=0, ) # Add to metrics at epoch 0 metrics["epoch"].append(int(0)) metrics["val_loss"].append(loss) metrics["val_iou"].append(iou) metrics["val_acc"].append(acc) metrics["train_loss"].append(None) metrics["train_iou"].append(None) metrics["train_acc"].append(None) initial_dice = dice best_model = model best_iou = iou patience = early_stopping epochs_without_improvement = 0 for epoch in tqdm( range(num_epochs), desc="Training Epochs", unit="epochs", ): epoch_loss = 0.0 num_batches = 0 model.train() # Ensure model is in training mode progress_bar = tqdm( train_dataloader, desc=f"Training Epoch {epoch + 1}", leave=True, unit="batches", ) for idx, batch in enumerate(progress_bar): # Get the inputs pixel_values = batch["pixel_values"].to(device) labels = batch["labels"].to(device) # Zero the parameter gradients optimizer.zero_grad() # Forward + backward + optimize outputs = model(pixel_values=pixel_values, labels=labels) loss, logits = outputs.loss, outputs.logits loss.backward() optimizer.step() epoch_loss += loss.item() num_batches += 1 # Evaluate training batch with torch.no_grad(): upsampled_logits = nn.functional.interpolate( logits, size=labels.shape[-2:], mode="bilinear", align_corners=False, ) predicted = upsampled_logits.argmax(dim=1) # Store for metric calculation pred_np = predicted.detach().cpu().numpy() ref_np = labels.detach().cpu().numpy() # Note: metric expects predictions + labels as numpy arrays metric.add_batch( predictions=pred_np, references=ref_np, ) train_metrics = metric.compute( num_labels=len(id2label), ignore_index=10, reduce_labels=False, ) train_loss = epoch_loss / num_batches if num_batches else 0.0 # Validation if (epoch + 1) % validate_every == 0: ( val_loss, val_iou, val_per_class_iou, val_acc, val_per_class_acc, val_dice, val_dice_per_class, ) = validate_model( model=model, dataloader=valid_dataloader, device=device, id2label=id2label, epoch=epoch + 1, ) # Record metrics metrics["epoch"].append(int(epoch + 1)) metrics["train_loss"].append(train_loss) metrics["train_iou"].append(train_metrics["mean_iou"]) metrics["train_acc"].append(train_metrics["mean_accuracy"]) metrics["val_loss"].append(val_loss) metrics["val_iou"].append(val_iou) metrics["val_acc"].append(val_acc) # Save the best model if val_iou > best_iou: best_model = model best_iou = val_iou epochs_without_improvement = 0 else: epochs_without_improvement += 1 if epochs_without_improvement >= patience: tqdm.write( f"Early stopping after {patience} epochs with no improvement", ) break return best_model, metrics, initial_dice def extract_model_zip(model_zip_path): """Extract model zip file and return the model directory.""" if not os.path.exists(model_zip_path): raise FileNotFoundError(f"Model zip file not found: {model_zip_path}") with zipfile.ZipFile(model_zip_path, "r") as zip_ref: extract_dir = os.path.join(os.path.dirname(model_zip_path), "output") zip_ref.extractall(extract_dir) # Check nested folder if len(os.listdir(extract_dir)) == 1: return os.path.join(extract_dir, os.listdir(extract_dir)[0]) else: return extract_dir def train_model( data_dir, base_model_zip, image_width, image_height, batch_size, data_percent, num_epochs, learning_rate, early_stopping, validate_every, ): model_path = extract_model_zip(base_model_zip) # Load model and labels model, id2label, id2color = load_model_and_labels(data_dir, model_path) # Create datasets and dataloaders train_dataloader, valid_dataloader = create_datasets_and_dataloaders( image_width, image_height, data_dir, batch_size, data_percent, ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # Train the model best_model, metrics, initial_dice = run_training( model, device, train_dataloader, valid_dataloader, id2label, num_epochs, learning_rate, early_stopping, validate_every, ) # Final validation ( loss, iou, per_class_iou, acc, per_class_acc, dice, dice_per_class, ) = validate_model( model=best_model, dataloader=valid_dataloader, device=device, id2label=id2label, calc_dice=True, epoch=0, ) final_dice = dice return best_model, metrics, [initial_dice, final_dice]