Scripts/inference.py

"""
Kurdish Handwritten Paragraph Recognition - Inference Script

Usage:
    # Single image
    python inference.py --image sample.tif --model_path model.safetensors --vocab_path vocab.json

    # Directory of images
    python inference.py --image_dir ./test_images --model_path model.safetensors --vocab_path vocab.json

    # With .pth checkpoint
    python inference.py --image sample.tif --model_path finetuned_model.pth --vocab_path vocab.json

    # KHATT Arabic model (different vocab)
    python inference.py --image arabic_sample.tif --model_path khatt_model.safetensors \
                        --vocab_path khatt_vocab.json
"""

import os
import glob
import json
import math
import time
import argparse
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.models as models


# ===============================
# Argument Parser
# ===============================

def parse_args():
    parser = argparse.ArgumentParser(
        description="Kurdish Handwritten Paragraph Recognition - Inference")

    # Input
    parser.add_argument("--image", type=str, default=None,
                        help="Path to a single paragraph image")
    parser.add_argument("--image_dir", type=str, default=None,
                        help="Directory of paragraph images to process")

    # Model and vocabulary
    parser.add_argument("--model_path", type=str, required=True,
                        help="Path to model weights (.pth or .safetensors)")
    parser.add_argument("--vocab_path", type=str, required=True,
                        help="Path to vocabulary JSON file (vocab.json)")
    parser.add_argument("--config_path", type=str, default=None,
                        help="Path to config.json (auto-loads architecture settings)")

    # Image dimensions
    parser.add_argument("--img_height", type=int, default=600)
    parser.add_argument("--img_width", type=int, default=1235)

    # Model architecture (overridden by config.json if provided)
    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=6)
    parser.add_argument("--num_heads", type=int, default=8)
    parser.add_argument("--ff_dim", type=int, default=2048)
    parser.add_argument("--max_seq_len", type=int, default=555)
    parser.add_argument("--use_upsample", action="store_true", default=True)
    parser.add_argument("--no_upsample", action="store_true")

    # Output
    parser.add_argument("--output_file", type=str, default=None,
                        help="Save predictions to text file")
    parser.add_argument("--show_timing", action="store_true",
                        help="Show per-image inference time")

    # Device
    parser.add_argument("--device", type=str, default=None,
                        help="Device (cuda/cpu, auto-detected if not set)")

    return parser.parse_args()


# ===============================
# Vocabulary
# ===============================

PAD_TOKEN = 0
SOS_TOKEN = 1
EOS_TOKEN = 2


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:
        mapping = vocab_data["char_to_idx"]
        char_list = [None] * len(mapping)
        for char, idx in mapping.items():
            char_list[idx] = char
    else:
        raise ValueError("Vocabulary JSON must contain 'vocab_list' or 'char_to_idx'")

    idx_to_char = {idx: char for idx, char in enumerate(char_list)}
    return char_list, idx_to_char


def decode_output(tensor, idx_to_char):
    """Convert tensor of character indices to 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


# ===============================
# Positional Encodings
# ===============================

class PositionalEncoding2D(nn.Module):
    """2D sinusoidal positional encoding for visual feature maps."""

    def __init__(self, d_model, max_h=100, max_w=300):
        super().__init__()
        pe = torch.zeros(max_h, max_w, d_model)
        d_half = d_model // 2

        pos_h = torch.arange(0, max_h, dtype=torch.float).unsqueeze(1)
        div_h = torch.exp(torch.arange(0, d_half, 2).float() * (-math.log(10000.0) / d_half))
        pe_h = torch.zeros(max_h, d_half)
        pe_h[:, 0::2] = torch.sin(pos_h * div_h)
        pe_h[:, 1::2] = torch.cos(pos_h * div_h)

        pos_w = torch.arange(0, max_w, dtype=torch.float).unsqueeze(1)
        div_w = torch.exp(torch.arange(0, d_half, 2).float() * (-math.log(10000.0) / d_half))
        pe_w = torch.zeros(max_w, d_half)
        pe_w[:, 0::2] = torch.sin(pos_w * div_w)
        pe_w[:, 1::2] = torch.cos(pos_w * div_w)

        for h in range(max_h):
            for w in range(max_w):
                pe[h, w, :d_half] = pe_h[h]
                pe[h, w, d_half:] = pe_w[w]

        self.register_buffer('pe', pe)

    def forward(self, x, height, width):
        _, seq_len, d_model = x.shape
        pe_2d = self.pe[:height, :width, :].reshape(height * width, d_model)
        if seq_len <= pe_2d.size(0):
            pe_2d = pe_2d[:seq_len]
        else:
            pad = torch.zeros(seq_len - pe_2d.size(0), d_model, device=x.device)
            pe_2d = torch.cat([pe_2d, pad], dim=0)
        return x + pe_2d.unsqueeze(0)


class PositionalEncoding1D(nn.Module):
    """1D sinusoidal positional encoding for decoder sequences."""

    def __init__(self, d_model, max_len=1000):
        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)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]


# ===============================
# CNN Feature Extractor
# ===============================

class CNNFeatureExtractor(nn.Module):
    """DenseNet-121 backbone with optional horizontal upsampling."""

    def __init__(self, output_dim=256, use_upsample=True):
        super().__init__()
        densenet = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
        self.features = densenet.features
        backbone_channels = 1024

        if use_upsample:
            self.upsample = nn.Sequential(
                nn.ConvTranspose2d(backbone_channels, 512,
                                   kernel_size=(1, 4), stride=(1, 2), padding=(0, 1)),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True))
            adapt_in = 512
        else:
            self.upsample = None
            adapt_in = backbone_channels

        self.adaptation = nn.Sequential(
            nn.Conv2d(adapt_in, output_dim, kernel_size=1),
            nn.BatchNorm2d(output_dim),
            nn.ReLU(inplace=True))

    def forward(self, x):
        features = F.relu(self.features(x), inplace=True)
        if self.upsample is not None:
            features = self.upsample(features)
        features = self.adaptation(features)
        b, c, h, w = features.shape
        return features.view(b, c, h * w).permute(0, 2, 1), h, w


# ===============================
# Transformer OCR Model
# ===============================

class TransformerOCRParagraphModel(nn.Module):
    """DenseNet121-Transformer for end-to-end paragraph recognition."""

    def __init__(self, vocab_size, hidden_size=256, nhead=8,
                 num_encoder_layers=3, num_decoder_layers=6,
                 dim_feedforward=2048, dropout=0.0,
                 use_upsample=True, max_seq_len=555):
        super().__init__()

        self.max_seq_len = max_seq_len
        self.vocab_size = vocab_size

        self.feature_extractor = CNNFeatureExtractor(
            output_dim=hidden_size, use_upsample=use_upsample)

        self.pos_encoder_2d = PositionalEncoding2D(hidden_size)
        self.pos_decoder_1d = PositionalEncoding1D(hidden_size, max_len=max_seq_len)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size, nhead=nhead,
            dim_feedforward=dim_feedforward, dropout=dropout,
            batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_encoder_layers)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_size, nhead=nhead,
            dim_feedforward=dim_feedforward, dropout=dropout,
            batch_first=True)
        self.transformer_decoder = nn.TransformerDecoder(
            decoder_layer, num_layers=num_decoder_layers)

        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.output_projection = nn.Linear(hidden_size, vocab_size)

    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 generate(self, img, max_length=None):
        """Auto-regressive greedy generation for a single image."""
        if max_length is None:
            max_length = self.max_seq_len

        self.eval()
        with torch.no_grad():
            if img.dim() == 3:
                img = img.unsqueeze(0)

            memory, feat_h, feat_w = self.feature_extractor(img)
            memory = self.pos_encoder_2d(memory, feat_h, feat_w)
            memory = self.transformer_encoder(memory)

            ys = torch.ones(1, 1).fill_(SOS_TOKEN).long().to(img.device)

            for _ in range(max_length - 1):
                tgt_embedded = self.pos_decoder_1d(self.token_embedding(ys))
                tgt_mask = self._generate_square_subsequent_mask(ys.size(1)).to(img.device)
                out = self.transformer_decoder(tgt_embedded, memory, 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 == EOS_TOKEN:
                    break

        return ys[0]

    def generate_batch(self, imgs, max_length=None):
        """Auto-regressive greedy batch generation."""
        if max_length is None:
            max_length = self.max_seq_len

        self.eval()
        batch_size = imgs.size(0)

        with torch.no_grad():
            memory, feat_h, feat_w = self.feature_extractor(imgs)
            memory = self.pos_encoder_2d(memory, feat_h, feat_w)
            memory = self.transformer_encoder(memory)

            ys = torch.ones(batch_size, 1).fill_(SOS_TOKEN).long().to(imgs.device)
            finished = torch.zeros(batch_size, dtype=torch.bool, device=imgs.device)

            for _ in range(max_length - 1):
                tgt_embedded = self.pos_decoder_1d(self.token_embedding(ys))
                tgt_mask = self._generate_square_subsequent_mask(ys.size(1)).to(imgs.device)
                out = self.transformer_decoder(tgt_embedded, memory, tgt_mask=tgt_mask)
                out = self.output_projection(out)

                next_tokens = out[:, -1].argmax(dim=-1)
                next_tokens[finished] = PAD_TOKEN
                ys = torch.cat([ys, next_tokens.unsqueeze(1)], dim=1)
                finished = finished | (next_tokens == EOS_TOKEN)
                if finished.all():
                    break

        return ys


# ===============================
# Image Preprocessing
# ===============================

def preprocess_image(image_path, img_height, img_width):
    """Load and preprocess a paragraph image.
    Aspect-ratio-preserving resize, right-aligned on white canvas for RTL."""
    image = Image.open(image_path).convert("RGB")
    orig_w, orig_h = image.size

    scale = min(img_width / orig_w, img_height / orig_h)
    new_w = int(orig_w * scale)
    new_h = int(orig_h * scale)
    image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)

    canvas = Image.new("RGB", (img_width, img_height), color=(255, 255, 255))
    x_offset = img_width - new_w  # Right-align for RTL
    canvas.paste(image, (x_offset, 0))

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    return transform(canvas)


# ===============================
# Config Loader
# ===============================

def load_config(config_path):
    """Load architecture settings from config.json."""
    with open(config_path, "r", encoding="utf-8") as f:
        return json.load(f)


# ===============================
# Main
# ===============================

def main():
    args = parse_args()

    if args.image is None and args.image_dir is None:
        print("Error: Provide --image or --image_dir")
        return

    # Device
    if args.device:
        device = torch.device(args.device)
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Load config if provided (overrides CLI args)
    if args.config_path and os.path.exists(args.config_path):
        config = load_config(args.config_path)
        print(f"Loaded config from: {args.config_path}")
        args.hidden_size = config.get("hidden_size", args.hidden_size)
        args.encoder_layers = config.get("num_encoder_layers", args.encoder_layers)
        args.decoder_layers = config.get("num_decoder_layers", args.decoder_layers)
        args.num_heads = config.get("num_attention_heads", args.num_heads)
        args.ff_dim = config.get("feed_forward_dim", args.ff_dim)
        args.max_seq_len = config.get("max_sequence_length", args.max_seq_len)
        args.img_height = config.get("image_height", args.img_height)
        args.img_width = config.get("image_width", args.img_width)
        if "use_upsample" in config:
            args.use_upsample = config["use_upsample"]
            args.no_upsample = not config["use_upsample"]

    use_upsample = args.use_upsample and not args.no_upsample

    # Vocabulary
    char_list, idx_to_char = load_vocabulary(args.vocab_path)
    vocab_size = len(char_list)
    print(f"Vocabulary: {vocab_size} tokens")

    # Model
    model = TransformerOCRParagraphModel(
        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,
        use_upsample=use_upsample,
        max_seq_len=args.max_seq_len
    ).to(device)

    # Load weights
    print(f"Loading weights: {args.model_path}")
    if args.model_path.endswith(".safetensors"):
        from safetensors.torch import load_file
        state_dict = load_file(args.model_path)
    else:
        checkpoint = torch.load(args.model_path, map_location=device)
        state_dict = checkpoint.get("model_state_dict", checkpoint)

    # Handle PE size mismatches
    model_state = model.state_dict()
    filtered = {}
    for key, value in state_dict.items():
        if key in model_state:
            if value.shape == model_state[key].shape:
                filtered[key] = value
    model.load_state_dict(filtered, strict=False)

    model.eval()
    total_params = sum(p.numel() for n, p in model.named_parameters() if '.pe' not in n)
    print(f"Model loaded: {total_params:,} parameters")
    print(f"Upsample: {'ON' if use_upsample else 'OFF'}")
    print(f"Image size: {args.img_height} x {args.img_width}")
    print(f"Max sequence length: {args.max_seq_len}")

    # Collect images
    image_paths = []
    if args.image:
        image_paths = [args.image]
    elif args.image_dir:
        for ext in ("*.tif", "*.tiff", "*.png", "*.jpg", "*.jpeg", "*.bmp"):
            image_paths.extend(glob.glob(os.path.join(args.image_dir, ext)))
            image_paths.extend(glob.glob(os.path.join(args.image_dir, ext.upper())))
        image_paths = sorted(list(set(image_paths)))

    if not image_paths:
        print("No images found.")
        return

    print(f"\nProcessing {len(image_paths)} image(s)...\n")

    # Output file
    out_file = None
    if args.output_file:
        out_file = open(args.output_file, "w", encoding="utf-8")

    total_time = 0

    for img_path in image_paths:
        filename = os.path.basename(img_path)

        # Preprocess
        tensor = preprocess_image(img_path, args.img_height, args.img_width).to(device)

        # Inference with timing
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        start = time.perf_counter()

        output = model.generate(tensor)

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - start
        total_time += elapsed

        # Decode
        text = decode_output(output, idx_to_char)
        lines = text.split('\n')

        # Display
        print(f"{'='*60}")
        print(f"File: {filename}")
        if args.show_timing:
            print(f"Time: {elapsed*1000:.1f} ms")
        print(f"Lines detected: {len(lines)}")
        print(f"{'─'*60}")
        for i, line in enumerate(lines):
            print(f"  Line {i+1}: {line}")
        print()

        # Save to file
        if out_file:
            out_file.write(f"# {filename}\n")
            out_file.write(text + "\n\n")

    # Summary
    print(f"{'='*60}")
    print(f"Done. {len(image_paths)} image(s) processed.")
    if args.show_timing:
        avg_ms = (total_time / len(image_paths)) * 1000
        print(f"Average inference: {avg_ms:.1f} ms/image")
        print(f"Total time: {total_time:.2f} s")

    if out_file:
        out_file.close()
        print(f"Predictions saved to: {args.output_file}")


if __name__ == "__main__":
    main()