Upload Scripts/train.py with huggingface_hub

Scripts/train.py

@@ -0,0 +1,718 @@
+"""
+Kurdish Handwritten Line Recognition - Training Script
+DenseNet121-Transformer Architecture with Constrained Synthetic Line Generation
+
+Usage:
+    python train.py --data_dir ./data/DASTNUS --vocab_path ./vocab.json
+    python train.py --data_dir ./data/DASTNUS --vocab_path ./vocab.json --use_synthetic --use_writer_mixing
+"""
+
+import os
+import glob
+import time
+import argparse
+import json
+import math
+import random
+import re
+import numpy as np
+from PIL import Image, ImageFilter
+from datetime import datetime
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.utils.data as data
+from torch.utils.data import ConcatDataset
+import torchvision.transforms as transforms
+import torchvision.models as models
+from torchvision.transforms import InterpolationMode
+from tqdm import tqdm
+
+# ===============================
+# Argument Parser
+# ===============================
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Kurdish Handwritten Line Recognition Training")
+
+    # Data paths
+    parser.add_argument("--data_dir", type=str, required=True,
+                        help="Root directory of DASTNUS dataset")
+    parser.add_argument("--vocab_path", type=str, required=True,
+                        help="Path to vocabulary JSON file (vocab.json)")
+    parser.add_argument("--synthetic_dir", type=str, default=None,
+                        help="Directory containing synthetic handwritten lines")
+    parser.add_argument("--fixed_lines_dir", type=str, default=None,
+                        help="Directory containing fixed-content handwritten lines")
+
+    # Data options
+    parser.add_argument("--use_synthetic", action="store_true",
+                        help="Include synthetic handwritten lines in training")
+    parser.add_argument("--use_writer_mixing", action="store_true",
+                        help="Include fixed-content lines from random writers")
+    parser.add_argument("--num_writers", type=int, default=50,
+                        help="Number of writers to randomly select for mixing")
+
+    # Image dimensions
+    parser.add_argument("--img_height", type=int, default=96)
+    parser.add_argument("--img_width", type=int, default=1235)
+
+    # Training hyperparameters
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--num_epochs", type=int, default=80)
+    parser.add_argument("--learning_rate", type=float, default=5e-4)
+    parser.add_argument("--grad_clip", type=float, default=5.0)
+    parser.add_argument("--weight_decay", type=float, default=1e-4)
+    parser.add_argument("--seed", type=int, default=42)
+
+    # Model parameters
+    parser.add_argument("--hidden_size", type=int, default=256)
+    parser.add_argument("--encoder_layers", type=int, default=3)
+    parser.add_argument("--decoder_layers", type=int, default=3)
+    parser.add_argument("--num_heads", type=int, default=8)
+    parser.add_argument("--dropout", type=float, default=0.4)
+    parser.add_argument("--ff_dim", type=int, default=1024)
+
+    # Early stopping
+    parser.add_argument("--patience", type=int, default=10)
+
+    # Augmentation
+    parser.add_argument("--no_aug", action="store_true",
+                        help="Disable adaptive augmentation")
+
+    # Output
+    parser.add_argument("--output_dir", type=str, default="./output",
+                        help="Directory to save models and logs")
+
+    return parser.parse_args()
+
+# ===============================
+# Vocabulary Loader
+# ===============================
+
+def load_vocabulary(vocab_path):
+    """Load vocabulary from JSON file"""
+    with open(vocab_path, "r", encoding="utf-8") as f:
+        vocab_data = json.load(f)
+
+    if "vocab_list" in vocab_data:
+        char_list = vocab_data["vocab_list"]
+    elif "char_to_idx" in vocab_data:
+        char_to_idx = vocab_data["char_to_idx"]
+        char_list = [None] * len(char_to_idx)
+        for char, idx in char_to_idx.items():
+            char_list[idx] = char
+    else:
+        raise ValueError("Vocabulary JSON must contain 'vocab_list' or 'char_to_idx'")
+
+    char_to_idx = {char: idx for idx, char in enumerate(char_list)}
+    idx_to_char = {idx: char for idx, char in enumerate(char_list)}
+
+    PAD_token = 0
+    SOS_token = 1
+    EOS_token = 2
+
+    return char_list, char_to_idx, idx_to_char, PAD_token, SOS_token, EOS_token
+
+# ===============================
+# Helper Functions
+# ===============================
+
+def tensor_to_text(tensor, idx_to_char, PAD_token, SOS_token, EOS_token):
+    """Convert a tensor of character indices to readable text"""
+    if isinstance(tensor, torch.Tensor):
+        tensor = tensor.cpu().tolist()
+    text = ""
+    for idx in tensor:
+        if idx == PAD_token or idx == SOS_token:
+            continue
+        if idx == EOS_token:
+            break
+        if idx in idx_to_char:
+            text += idx_to_char[idx]
+    return text
+
+def extract_writer_id(filename):
+    """Extract writer ID from filename (e.g., DNDK00002_2_1.tif -> 2)"""
+    basename = os.path.basename(filename)
+    match = re.match(r"DNDK(\d+)", basename)
+    if match:
+        return int(match.group(1))
+    return None
+
+def get_unique_writers(directory):
+    """Get all unique writer IDs from a directory"""
+    image_files = glob.glob(os.path.join(directory, "*.tif"))
+    writer_ids = set()
+    for f in image_files:
+        wid = extract_writer_id(f)
+        if wid is not None:
+            writer_ids.add(wid)
+    return sorted(list(writer_ids))
+
+def filter_files_by_writers(directory, selected_writers):
+    """Filter image files to only include those from selected writers"""
+    all_files = glob.glob(os.path.join(directory, "*.tif"))
+    return [f for f in all_files if extract_writer_id(f) in selected_writers]
+
+# ===============================
+# Dataset Class
+# ===============================
+
+# Global variables for adaptive augmentation
+current_epoch = 0
+num_epochs_global = 80
+overfitting_detected = False
+validation_loss_history = []
+training_loss_history = []
+
+class KurdishLineDataset(data.Dataset):
+    def __init__(self, root_dir=None, transform=None, max_samples=None,
+                 dataset_name="", image_files=None, img_height=96, img_width=1235,
+                 char_to_idx=None, SOS_token=1, EOS_token=2):
+        self.transform = transform
+        self.dataset_name = dataset_name
+        self.img_height = img_height
+        self.img_width = img_width
+        self.char_to_idx = char_to_idx
+        self.SOS_token = SOS_token
+        self.EOS_token = EOS_token
+
+        if image_files is not None:
+            self.image_files = image_files
+        else:
+            self.image_files = glob.glob(os.path.join(root_dir, "*.tif"))
+
+        if max_samples and max_samples < len(self.image_files):
+            self.image_files = self.image_files[:max_samples]
+
+        print(f"Loaded {len(self.image_files)} images for {dataset_name}")
+
+    def __len__(self):
+        return len(self.image_files)
+
+    def __getitem__(self, idx):
+        img_path = self.image_files[idx]
+        label_path = os.path.splitext(img_path)[0] + ".txt"
+
+        image = Image.open(img_path).convert("RGB")
+        orig_width, orig_height = image.size
+        aspect_ratio = orig_width / orig_height
+
+        new_height = self.img_height
+        new_width = min(int(new_height * aspect_ratio), self.img_width)
+        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+        target_img = Image.new("RGB", (self.img_width, self.img_height), color=(255, 255, 255))
+        target_img.paste(image, (0, 0))
+
+        if self.transform:
+            target_img = self.transform(target_img)
+
+        try:
+            with open(label_path, "r", encoding="utf-8") as f:
+                text = f.readline().strip()
+        except UnicodeDecodeError:
+            with open(label_path, "r", encoding="utf-8-sig") as f:
+                text = f.readline().strip()
+
+        indices = ([self.SOS_token] +
+                   [self.char_to_idx.get(char, self.char_to_idx.get(" ", 0)) for char in text] +
+                   [self.EOS_token])
+
+        target = torch.LongTensor(indices)
+        target_length = len(indices)
+
+        return target_img, target, target_length, text
+
+def collate_fn(batch):
+    """Custom collate function for padding sequences"""
+    batch.sort(key=lambda x: x[2], reverse=True)
+    images, targets, lengths, original_texts = zip(*batch)
+
+    images = torch.stack(images, 0)
+    max_length = max(lengths)
+
+    padded_targets = torch.zeros(len(targets), max_length).long()
+    for i, target in enumerate(targets):
+        padded_targets[i, :lengths[i]] = target[:lengths[i]]
+
+    lengths = torch.LongTensor(lengths)
+    return images, padded_targets, lengths, original_texts
+
+# ===============================
+# Adaptive Augmentation
+# ===============================
+
+class AdaptiveStrokeWidthJitter:
+    def __init__(self, base_p=0.2, max_p=0.6, base_kernel=3, max_kernel=5):
+        self.base_p, self.max_p = base_p, max_p
+        self.base_kernel, self.max_kernel = base_kernel, max_kernel
+
+    def __call__(self, img):
+        progress = min(current_epoch / num_epochs_global, 1.0)
+        factor = 1.5 if overfitting_detected else 1.0
+        p = min(self.base_p + (self.max_p - self.base_p) * progress * factor, self.max_p)
+        kernel = self.base_kernel + int(2 * progress)
+        if kernel % 2 == 0:
+            kernel += 1
+        kernel = min(kernel, self.max_kernel)
+        if random.random() < p:
+            if random.random() < 0.5:
+                return img.filter(ImageFilter.MinFilter(kernel))
+            return img.filter(ImageFilter.MaxFilter(kernel))
+        return img
+
+class AdaptiveGaussianNoise:
+    def __init__(self, base_std=(0.0, 0.01), max_std=(0.0, 0.03), base_p=0.3, max_p=0.7):
+        self.base_std, self.max_std = base_std, max_std
+        self.base_p, self.max_p = base_p, max_p
+
+    def __call__(self, tensor):
+        progress = min(current_epoch / num_epochs_global, 1.0)
+        factor = 1.5 if overfitting_detected else 1.0
+        p = min(self.base_p + (self.max_p - self.base_p) * progress * factor, self.max_p)
+        std_high = min(self.base_std[1] + (self.max_std[1] - self.base_std[1]) * progress * factor,
+                       self.max_std[1])
+        if torch.rand(1).item() < p:
+            noise = torch.randn_like(tensor) * random.uniform(self.base_std[0], std_high)
+            tensor = torch.clamp(tensor + noise, 0.0, 1.0)
+        return tensor
+
+def build_adaptive_train_transform():
+    class AdaptiveTransform:
+        def __call__(self, img):
+            progress = min(current_epoch / num_epochs_global, 1.0)
+            factor = 1.3 if overfitting_detected else 1.0
+
+            if random.random() < min(0.6 + 0.35 * progress * factor, 0.95):
+                b = min(0.1 + 0.2 * progress * factor, 0.3)
+                img = transforms.ColorJitter(brightness=b, contrast=b)(img)
+
+            if random.random() < min(0.7 + 0.25 * progress * factor, 0.95):
+                deg = min(1 + 4 * progress * factor, 5)
+                shear = min(3 + 7 * progress * factor, 10)
+                img = transforms.RandomAffine(
+                    degrees=deg,
+                    translate=(min(0.01 + 0.02 * progress, 0.03),
+                               min(0.03 + 0.05 * progress, 0.08)),
+                    scale=(max(1 - 0.02 - 0.08 * progress, 0.90),
+                           min(1 + 0.02 + 0.08 * progress, 1.10)),
+                    shear=(-shear, shear),
+                    interpolation=InterpolationMode.BILINEAR, fill=255)(img)
+
+            if random.random() < min(0.1 + 0.4 * progress * factor, 0.5):
+                dist = min(0.02 + 0.06 * progress * factor, 0.08)
+                img = transforms.RandomPerspective(
+                    distortion_scale=dist, p=1.0,
+                    interpolation=InterpolationMode.BILINEAR, fill=255)(img)
+
+            if random.random() < min(0.15 + 0.2 * progress, 0.35):
+                img = transforms.GaussianBlur(
+                    kernel_size=3, sigma=(0.1, min(0.5 + 0.5 * progress, 1.0)))(img)
+
+            img = AdaptiveStrokeWidthJitter()(img)
+            img = transforms.ToTensor()(img)
+            img = AdaptiveGaussianNoise()(img)
+
+            if random.random() < min(0.1 + 0.3 * progress * factor, 0.4):
+                img = transforms.RandomErasing(
+                    p=1.0, scale=(0.01, min(0.01 + 0.04 * progress, 0.05)),
+                    ratio=(0.3, 3.3), value="random")(img)
+
+            img = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))(img)
+            return img
+
+    return AdaptiveTransform()
+
+def build_eval_transform():
+    return transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+    ])
+
+# ===============================
+# Model Architecture
+# ===============================
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        return x + self.pe[:x.size(0), :]
+
+class DenseNetFeatureExtractor(nn.Module):
+    def __init__(self, output_dim=256):
+        super().__init__()
+        densenet = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
+        self.features = nn.Sequential(*list(densenet.children())[:-1])
+        self.adapt = nn.Conv2d(1024, output_dim, kernel_size=1)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.adapt(x)
+        x = nn.functional.adaptive_avg_pool2d(x, (1, None))
+        x = x.squeeze(2)
+        return x.permute(0, 2, 1)
+
+class TransformerOCRModel(nn.Module):
+    def __init__(self, vocab_size, hidden_size=256, nhead=8,
+                 num_encoder_layers=3, num_decoder_layers=3,
+                 dim_feedforward=1024, dropout=0.4,
+                 PAD_token=0, SOS_token=1, EOS_token=2, max_seq_len=150):
+        super().__init__()
+        self.feature_extractor = DenseNetFeatureExtractor(output_dim=hidden_size)
+        self.pos_encoder = PositionalEncoding(hidden_size)
+        self.transformer = nn.Transformer(
+            d_model=hidden_size, nhead=nhead,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout, batch_first=True)
+        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
+        self.output_projection = nn.Linear(hidden_size, vocab_size)
+        self.hidden_size = hidden_size
+        self.vocab_size = vocab_size
+        self.PAD_token = PAD_token
+        self.SOS_token = SOS_token
+        self.EOS_token = EOS_token
+        self.max_seq_len = max_seq_len
+        self._init_parameters()
+
+    def _init_parameters(self):
+        nn.init.xavier_uniform_(self.token_embedding.weight)
+        nn.init.xavier_uniform_(self.output_projection.weight)
+
+    def _generate_square_subsequent_mask(self, sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        return mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, 0.0)
+
+    def _apply_pos_encoding(self, x):
+        x = x.permute(1, 0, 2)
+        x = self.pos_encoder(x)
+        return x.permute(1, 0, 2)
+
+    def forward(self, src, tgt):
+        memory = self._apply_pos_encoding(self.feature_extractor(src))
+        tgt_input = tgt[:, :-1]
+        tgt_embedded = self._apply_pos_encoding(self.token_embedding(tgt_input))
+        tgt_mask = self._generate_square_subsequent_mask(tgt_embedded.size(1)).to(src.device)
+        output = self.transformer(
+            src=memory, tgt=tgt_embedded, tgt_mask=tgt_mask,
+            src_is_causal=False, tgt_is_causal=True)
+        return self.output_projection(output)
+
+    def generate(self, img):
+        """Generate text from a single image"""
+        was_training = self.training
+        self.eval()
+        with torch.no_grad():
+            if img.dim() == 3:
+                img = img.unsqueeze(0)
+            memory = self._apply_pos_encoding(self.feature_extractor(img))
+            ys = torch.ones(1, 1).fill_(self.SOS_token).long().to(img.device)
+
+            for _ in range(self.max_seq_len - 1):
+                tgt_embedded = self._apply_pos_encoding(self.token_embedding(ys))
+                tgt_mask = self._generate_square_subsequent_mask(ys.size(1)).to(img.device)
+                out = self.transformer(src=memory, tgt=tgt_embedded, tgt_mask=tgt_mask)
+                out = self.output_projection(out)
+                next_word = out[0, -1].argmax().item()
+                ys = torch.cat([ys, torch.ones(1, 1).long().fill_(next_word).to(img.device)], dim=1)
+                if next_word == self.EOS_token:
+                    break
+
+        if was_training:
+            self.train(True)
+        return ys[0]
+
+    def generate_batch(self, imgs):
+        """Generate text from a batch of images"""
+        self.eval()
+        batch_size = imgs.size(0)
+        with torch.no_grad():
+            memory = self._apply_pos_encoding(self.feature_extractor(imgs))
+            ys = torch.ones(batch_size, 1).fill_(self.SOS_token).long().to(imgs.device)
+            finished = torch.zeros(batch_size, dtype=torch.bool, device=imgs.device)
+
+            for _ in range(self.max_seq_len - 1):
+                tgt_embedded = self._apply_pos_encoding(self.token_embedding(ys))
+                tgt_mask = self._generate_square_subsequent_mask(ys.size(1)).to(imgs.device)
+                out = self.transformer(src=memory, tgt=tgt_embedded, tgt_mask=tgt_mask)
+                out = self.output_projection(out)
+                next_tokens = out[:, -1].argmax(dim=-1)
+                next_tokens[finished] = self.PAD_token
+                ys = torch.cat([ys, next_tokens.unsqueeze(1)], dim=1)
+                finished = finished | (next_tokens == self.EOS_token)
+                if finished.all():
+                    break
+        return ys
+
+# ===============================
+# Metrics
+# ===============================
+
+def levenshtein_distance(s1, s2):
+    if len(s1) < len(s2):
+        return levenshtein_distance(s2, s1)
+    if len(s2) == 0:
+        return len(s1)
+    prev = range(len(s2) + 1)
+    for c1 in s1:
+        curr = [prev[0] + 1]
+        for j, c2 in enumerate(s2):
+            curr.append(min(prev[j + 1] + 1, curr[j] + 1, prev[j] + (c1 != c2)))
+        prev = curr
+    return prev[-1]
+
+def calculate_cer(preds, targets):
+    total_dist, total_chars = 0, 0
+    for p, t in zip(preds, targets):
+        total_dist += levenshtein_distance(p, t)
+        total_chars += len(t)
+    return total_dist / max(1, total_chars)
+
+def calculate_wer(preds, targets):
+    total_dist, total_words = 0, 0
+    for p, t in zip(preds, targets):
+        total_dist += levenshtein_distance(p.split(), t.split())
+        total_words += len(t.split())
+    return total_dist / max(1, total_words)
+
+def evaluate_cer(model, dataloader, device, idx_to_char, PAD_token, SOS_token, EOS_token):
+    model.eval()
+    all_preds, all_targets = [], []
+    with torch.no_grad():
+        for images, _, _, texts in tqdm(dataloader, desc="Evaluating"):
+            images = images.to(device)
+            batch_output = model.generate_batch(images)
+            for seq in batch_output:
+                all_preds.append(tensor_to_text(seq, idx_to_char, PAD_token, SOS_token, EOS_token))
+            all_targets.extend(texts)
+    cer = calculate_cer(all_preds, all_targets)
+    return cer, all_preds, all_targets
+
+# ===============================
+# Early Stopping
+# ===============================
+
+class EarlyStopping:
+    def __init__(self, patience=10):
+        self.patience = patience
+        self.counter = 0
+        self.best_cer = float("inf")
+        self.early_stop = False
+
+    def __call__(self, val_cer, model, epoch, path):
+        if val_cer < self.best_cer:
+            self.best_cer = val_cer
+            self.counter = 0
+            torch.save({
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "val_cer": val_cer
+            }, path)
+            print(f"Model saved (Val CER: {val_cer:.4f})")
+        else:
+            self.counter += 1
+            print(f"Early stopping: {self.counter}/{self.patience}")
+            if self.counter >= self.patience:
+                self.early_stop = True
+                print("Early stopping triggered.")
+
+# ===============================
+# Training Loop
+# ===============================
+
+def train_epoch(model, dataloader, optimizer, criterion, device, scheduler, PAD_token):
+    model.train()
+    epoch_loss = 0
+    for images, targets, _, _ in tqdm(dataloader, desc="Training"):
+        images, targets = images.to(device), targets.to(device)
+        optimizer.zero_grad()
+        outputs = model(images, targets)
+        outputs = outputs.reshape(-1, outputs.shape[-1])
+        targets_flat = targets[:, 1:].reshape(-1)
+        loss = criterion(outputs, targets_flat)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        optimizer.step()
+        scheduler.step()
+        epoch_loss += loss.item()
+    return epoch_loss / len(dataloader)
+
+def evaluate_model(model, dataloader, criterion, device, PAD_token):
+    model.eval()
+    epoch_loss = 0
+    with torch.no_grad():
+        for images, targets, _, _ in dataloader:
+            images, targets = images.to(device), targets.to(device)
+            outputs = model(images, targets)
+            outputs = outputs.reshape(-1, outputs.shape[-1])
+            targets_flat = targets[:, 1:].reshape(-1)
+            loss = criterion(outputs, targets_flat)
+            epoch_loss += loss.item()
+    return epoch_loss / len(dataloader)
+
+# ===============================
+# Main
+# ===============================
+
+def main():
+    global current_epoch, num_epochs_global, overfitting_detected
+    global validation_loss_history, training_loss_history, args
+
+    args = parse_args()
+
+    # Set seeds
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    num_epochs_global = args.num_epochs
+
+    # Device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # Output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Load vocabulary
+    char_list, char_to_idx, idx_to_char, PAD_token, SOS_token, EOS_token = \
+        load_vocabulary(args.vocab_path)
+    vocab_size = len(char_list)
+    print(f"Vocabulary size: {vocab_size}")
+
+    # Transforms
+    train_transform = build_eval_transform() if args.no_aug else build_adaptive_train_transform()
+    eval_transform = build_eval_transform()
+
+    # Dataset common kwargs
+    ds_kwargs = dict(img_height=args.img_height, img_width=args.img_width,
+                     char_to_idx=char_to_idx, SOS_token=SOS_token, EOS_token=EOS_token)
+
+    # Build datasets
+    real_train_dir = os.path.join(args.data_dir, "Training")
+    real_val_dir = os.path.join(args.data_dir, "Validation")
+    real_test_dir = os.path.join(args.data_dir, "Testing")
+
+    train_datasets = [
+        KurdishLineDataset(real_train_dir, transform=train_transform,
+                           dataset_name="Real Training", **ds_kwargs)
+    ]
+
+    if args.use_synthetic and args.synthetic_dir:
+        syn_dir = os.path.join(args.synthetic_dir, "Training")
+        train_datasets.append(
+            KurdishLineDataset(syn_dir, transform=train_transform,
+                               dataset_name="Synthetic Training", **ds_kwargs))
+
+    if args.use_writer_mixing and args.fixed_lines_dir:
+        fix_dir = os.path.join(args.fixed_lines_dir, "Training")
+        all_writers = get_unique_writers(fix_dir)
+        selected = random.sample(all_writers, min(args.num_writers, len(all_writers)))
+        selected_files = filter_files_by_writers(fix_dir, set(selected))
+        train_datasets.append(
+            KurdishLineDataset(image_files=selected_files, transform=train_transform,
+                               dataset_name=f"Fixed {len(selected)} Writers", **ds_kwargs))
+
+    train_dataset = ConcatDataset(train_datasets) if len(train_datasets) > 1 else train_datasets[0]
+    val_dataset = KurdishLineDataset(real_val_dir, transform=eval_transform,
+                                     dataset_name="Validation", **ds_kwargs)
+    test_dataset = KurdishLineDataset(real_test_dir, transform=eval_transform,
+                                      dataset_name="Testing", **ds_kwargs)
+
+    print(f"\nTraining: {len(train_dataset)} | Validation: {len(val_dataset)} | Testing: {len(test_dataset)}")
+
+    # Data loaders
+    loader_kwargs = dict(num_workers=0, pin_memory=True, collate_fn=collate_fn)
+    train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **loader_kwargs)
+    val_loader = data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, **loader_kwargs)
+    test_loader = data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, **loader_kwargs)
+
+    # Model
+    model = TransformerOCRModel(
+        vocab_size=vocab_size, hidden_size=args.hidden_size,
+        nhead=args.num_heads, num_encoder_layers=args.encoder_layers,
+        num_decoder_layers=args.decoder_layers, dim_feedforward=args.ff_dim,
+        dropout=args.dropout, PAD_token=PAD_token, SOS_token=SOS_token,
+        EOS_token=EOS_token).to(device)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {total_params:,}")
+
+    # Optimizer and schedulers
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+    onecycle = optim.lr_scheduler.OneCycleLR(
+        optimizer, max_lr=args.learning_rate,
+        steps_per_epoch=len(train_loader), epochs=args.num_epochs, pct_start=0.1)
+    plateau = optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="min", factor=0.5, patience=2, min_lr=1e-6)
+
+    criterion = nn.CrossEntropyLoss(ignore_index=PAD_token)
+    early_stopping = EarlyStopping(patience=args.patience)
+    best_model_path = os.path.join(args.output_dir, "best_model.pth")
+
+    # Training
+    print(f"\nStarting training for {args.num_epochs} epochs...")
+    best_val_cer = float("inf")
+
+    for epoch in range(args.num_epochs):
+        current_epoch = epoch
+        start = time.time()
+
+        train_loss = train_epoch(model, train_loader, optimizer, criterion, device, onecycle, PAD_token)
+        val_loss = evaluate_model(model, val_loader, criterion, device, PAD_token)
+        val_cer, _, _ = evaluate_cer(model, val_loader, device, idx_to_char,
+                                      PAD_token, SOS_token, EOS_token)
+
+        # Overfitting detection
+        training_loss_history.append(train_loss)
+        validation_loss_history.append(val_loss)
+        if len(training_loss_history) >= 3:
+            overfitting_detected = (np.mean(validation_loss_history[-3:]) >
+                                    np.mean(training_loss_history[-3:]) * 1.2)
+
+        plateau.step(val_cer)
+        mins, secs = divmod(time.time() - start, 60)
+
+        print(f"\nEpoch {epoch + 1}/{args.num_epochs} ({mins:.0f}m {secs:.0f}s)")
+        print(f"  Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val CER: {val_cer:.4f}")
+
+        if val_cer < best_val_cer:
+            best_val_cer = val_cer
+
+        early_stopping(val_cer, model, epoch, best_model_path)
+        if early_stopping.early_stop:
+            break
+
+    # Final evaluation
+    print(f"\nBest validation CER: {best_val_cer:.4f}")
+    print("Loading best model for test evaluation...")
+    checkpoint = torch.load(best_model_path)
+    model.load_state_dict(checkpoint["model_state_dict"])
+
+    test_cer, test_preds, test_targets = evaluate_cer(
+        model, test_loader, device, idx_to_char, PAD_token, SOS_token, EOS_token)
+    test_wer = calculate_wer(test_preds, test_targets)
+
+    print(f"\nTest CER: {test_cer:.4f}")
+    print(f"Test WER: {test_wer:.4f}")
+    print(f"Test CRR: {(1 - test_cer) * 100:.2f}%")
+
+    for i in range(min(5, len(test_preds))):
+        print(f"\nSample {i + 1}:")
+        print(f"  Predicted: {test_preds[i]}")
+        print(f"  Actual:    {test_targets[i]}")
+
+if __name__ == "__main__":
+    main()