build_arc1.py

import os, random, torch, torch.nn.functional as F
from torch.utils.data import DataLoader, random_split, Dataset
import torch.nn as nn
from torch.cuda.amp import GradScaler
import matplotlib.pyplot as plt
import numpy as np
import logging
from tqdm import tqdm

# ============================================================
# Dataset
# ============================================================
class PlateDataset(Dataset):
    def __init__(self, csv_file, vocab, img_dir, transform=None, train=True):
        import pandas as pd
        from PIL import Image

        self.df = pd.read_csv(csv_file)
        self.vocab = vocab
        self.char2idx = {c: i for i, c in enumerate(vocab)}
        self.img_dir = img_dir
        self.transform = transform
        self.train = train

    def text_to_seq(self, text):
        return [self.char2idx['<sos>']] + \
               [self.char2idx[c] for c in text if c in self.char2idx] + \
               [self.char2idx['<eos>']]

    def seq_to_text(self, seq):
        return ''.join([self.vocab[i] for i in seq if i > 2])

    def __getitem__(self, idx):
        import cv2
        row = self.df.iloc[idx]
        img_path = os.path.join(self.img_dir, row['filename'])
        print(img_path)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (512, 96))  # support 96x512
        img = img.astype(np.float32) / 255.0
        img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
        img = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1)


        text = str(row['plate_text'])
        seq = torch.tensor(self.text_to_seq(text), dtype=torch.long)
        return img, seq, text

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


def collate_fn(batch):
    imgs, seqs, texts = zip(*batch)
    imgs = torch.stack(imgs, 0)
    lengths = [len(s) for s in seqs]
    max_len = max(lengths)
    padded = torch.zeros(len(seqs), max_len, dtype=torch.long)
    for i, s in enumerate(seqs):
        padded[i, :len(s)] = s
    return imgs, padded, texts


# ============================================================
# Model (simple CNN + Transformer decoder skeleton)
# ============================================================
class FastPlateOCR(nn.Module):
    def __init__(self, vocab_size, hidden=256, num_layers=4, nhead=8):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(3, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(128, hidden, 3, stride=2, padding=1), nn.ReLU(),
        )
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden, nhead=nhead, dim_feedforward=hidden*4,
            dropout=0.1, activation="relu", batch_first=True
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.embedding = nn.Embedding(vocab_size, hidden)
        self.fc = nn.Linear(hidden, vocab_size)

    def forward(self, imgs, tgt_inp):
        feats = self.cnn(imgs)               # [B,C,H,W]
        feats = feats.flatten(2).permute(0, 2, 1)  # [B,Seq,C]
        tgt = self.embedding(tgt_inp)        # [B,T,C]
        out = self.decoder(tgt, feats)       # [B,T,C]
        return self.fc(out)

    def greedy_decode(self, imgs, sos_id=1, eos_id=2, max_len=32, device="cpu"):
        self.eval()
        with torch.no_grad():
            feats = self.cnn(imgs).flatten(2).permute(0, 2, 1)
            ys = torch.ones(imgs.size(0), 1, dtype=torch.long, device=device) * sos_id
            for _ in range(max_len):
                out = self.decoder(self.embedding(ys), feats)
                prob = self.fc(out[:, -1])
                _, next_word = torch.max(prob, dim=1)
                ys = torch.cat([ys, next_word.unsqueeze(1)], dim=1)
                if next_word.item() == eos_id:
                    break
        return ys.squeeze(0).tolist()

    def beam_decode(self, imgs, beam_width=5, sos_id=1, eos_id=2, max_len=32, device="cpu"):
        self.eval()
        with torch.no_grad():
            feats = self.cnn(imgs).flatten(2).permute(0, 2, 1)
            sequences = [[torch.tensor([sos_id], device=device), 0.0]]
            for _ in range(max_len):
                all_candidates = []
                for seq, score in sequences:
                    if seq[-1].item() == eos_id:
                        all_candidates.append((seq, score))
                        continue
                    out = self.decoder(self.embedding(seq.unsqueeze(0)), feats)
                    prob = torch.log_softmax(self.fc(out[:, -1]), dim=-1)
                    topk = torch.topk(prob, beam_width)
                    for i in range(beam_width):
                        next_seq = torch.cat([seq, topk.indices[0][i].unsqueeze(0)])
                        all_candidates.append((next_seq, score + topk.values[0][i].item()))
                ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True)
                sequences = ordered[:beam_width]
            return sequences[0][0].tolist()


# ============================================================
# Validation
# ============================================================
def validate(model, val_loader, dataset, device, beam_width=5):
    model.eval()
    total_cer, total_acc, count = 0, 0, 0
    import editdistance
    with torch.no_grad():
        for imgs, seqs, texts in val_loader:
            imgs = imgs.to(device)
            pred_ids = model.beam_decode(imgs, beam_width=beam_width, device=device)
            pred_text = dataset.seq_to_text(pred_ids)
            gt_text = texts[0]

            # CER
            cer = editdistance.eval(pred_text, gt_text) / max(1, len(gt_text))
            total_cer += cer
            total_acc += (pred_text == gt_text)
            count += 1
    return total_cer / count, total_acc / count




def get_logger(log_file):
    logger = logging.getLogger("train_logger")
    logger.setLevel(logging.INFO)

    # 🔁 Remove all existing handlers to avoid stale configs
    while logger.handlers:
        handler = logger.handlers[0]
        logger.removeHandler(handler)
        handler.close()

    # 📁 File handler
    fh = logging.FileHandler(log_file, mode='a')
    fh.setLevel(logging.INFO)
    fh.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
    logger.addHandler(fh)

    # 🖥️ Console handler
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    ch.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
    logger.addHandler(ch)

    return logger


# ---------------------------
# Training function
# ---------------------------
def train_fastplateocr(csv_file, img_dir="src2/resized_plates",
                       epochs=20, batch_size=16, beam_width=5, lr=1e-4,
                       out_dir="results2"):

    charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    vocab = ['<pad>', '<sos>', '<eos>'] + list(charset)

    dataset = PlateDataset(csv_file, vocab, img_dir=img_dir, train=True)
    val_size = max(1, int(0.1 * len(dataset)))
    train_size = len(dataset) - val_size
    train_ds, val_ds = random_split(dataset, [train_size, val_size])

    train_loader = DataLoader(train_ds, batch_size=batch_size,
                              shuffle=True, collate_fn=collate_fn, num_workers=2)
    val_loader = DataLoader(val_ds, batch_size=1,
                            shuffle=False, collate_fn=collate_fn, num_workers=1)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = FastPlateOCR(vocab_size=len(vocab)).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    scaler = GradScaler()

    os.makedirs(out_dir, exist_ok=True)
    log_file = os.path.join(out_dir, "training_logs.txt")
    logger = get_logger(log_file)

    logs = {"train_losses": [], "val_cers": [], "val_accs": []}
    best_cer = float("inf")

    logger.info("🟢 Training started")

    for epoch in range(1, epochs+1):
        model.train()
        running_loss = 0.0

        for imgs, seqs, _ in tqdm(train_loader, desc=f"Epoch {epoch}/{epochs}", leave=False, dynamic_ncols=True):
            imgs, seqs = imgs.to(device), seqs.to(device)
            tgt_inp, tgt_out = seqs[:, :-1], seqs[:, 1:]

            optimizer.zero_grad()
            with torch.cuda.amp.autocast():
                logits = model(imgs, tgt_inp)
                loss = F.cross_entropy(
                    logits.reshape(-1, logits.size(-1)),
                    tgt_out.reshape(-1),
                    ignore_index=0
                )

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            running_loss += loss.item()

        scheduler.step()
        avg_loss = running_loss / len(train_loader)

        val_cer, val_acc = validate(model, val_loader, dataset, device, beam_width)

        logs["train_losses"].append(avg_loss)
        logs["val_cers"].append(val_cer)
        logs["val_accs"].append(val_acc)

        lr = scheduler.get_last_lr()[0]
        epoch_log = (
            f"Epoch {epoch}/{epochs} | "
            f"Train Loss: {avg_loss:.6f} | "
            f"Val CER: {val_cer:.6f} | "
            f"Val Acc: {val_acc:.6f} | "
            f"LR: {lr:.8f}"
        )

        # Log to console and file
        logger.info(epoch_log)
        for handler in logger.handlers:
            handler.flush()  # force flush to file

        # Save best model
        if val_cer < best_cer:
            best_cer = val_cer
            torch.save({"model": model.state_dict()}, os.path.join(out_dir, "best.pth"))
            logger.info(f"💾 Saved checkpoint (epoch {epoch}) — new best CER: {best_cer:.6f}")
            for handler in logger.handlers:
                handler.flush()
    # Save training curves
    plt.figure(figsize=(12,4))
    plt.subplot(1,3,1); plt.plot(logs["train_losses"], label="Train Loss"); plt.title("Train Loss"); plt.legend()
    plt.subplot(1,3,2); plt.plot(logs["val_cers"], label="Val CER"); plt.title("Val CER"); plt.legend()
    plt.subplot(1,3,3); plt.plot(logs["val_accs"], label="Val Acc"); plt.title("Val Accuracy"); plt.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, "training_curves.png"))
    plt.close()

    logger.info("✅ Training complete. Curves saved.")

    return model, vocab, logs

# ---------------------------
# Run
# ---------------------------
if __name__ == "__main__":
    train_fastplateocr("clean_plates.csv", img_dir="synthetic_plates1")