submission_code/ocr_helpers.py

from __future__ import annotations

import csv
import math
import re
import unicodedata
import zipfile
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Iterable, Sequence

import numpy as np
from PIL import Image


@dataclass
class ImageConfig:
    target_height: int | None = None
    target_width: int | None = None
    preserve_aspect_ratio: bool = True
    grayscale: bool = True
    normalize: bool = True
    binarize: bool = False


@dataclass
class TextConfig:
    unicode_form: str = "NFC"
    remove_tatweel: bool = True
    remove_diacritics: bool = False
    keep_punctuation: bool = True
    collapse_spaces: bool = True
    strip_text: bool = True


TATWEEL = "\u0640"
DIACRITIC_PATTERN = re.compile(r"[\u064b-\u065f\u0670]")
SPACE_PATTERN = re.compile(r"\s+")
ARABIC_PUNCTUATION = "،؛؟"
ASCII_PUNCTUATION = r"!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
PUNCTUATION_PATTERN = re.compile("[" + re.escape(ASCII_PUNCTUATION + ARABIC_PUNCTUATION) + "]")


AHCD_CLASS_LABELS = {
    1: "alef",
    2: "beh",
    3: "teh",
    4: "theh",
    5: "jeem",
    6: "hah",
    7: "khah",
    8: "dal",
    9: "thal",
    10: "reh",
    11: "zain",
    12: "seen",
    13: "sheen",
    14: "sad",
    15: "dad",
    16: "tah",
    17: "zah",
    18: "ain",
    19: "ghain",
    20: "feh",
    21: "qaf",
    22: "kaf",
    23: "lam",
    24: "meem",
    25: "noon",
    26: "heh",
    27: "waw",
    28: "yeh",
}

AHCD_FILE_PATTERN = re.compile(r"(?:train|test)/id_(?P<sample_id>\d+)_label_(?P<label>\d+)\.png$")


def project_root_from_notebook() -> Path:
    """Return the submission folder when called from any notebook in this package."""
    cwd = Path.cwd().resolve()
    if cwd.name == "notebooks" and cwd.parent.name == "submission":
        return cwd.parent
    if cwd.name == "submission":
        return cwd
    candidate = cwd / "submission"
    return candidate if candidate.exists() else cwd


def normalize_arabic_text(text: str, config: TextConfig | None = None) -> str:
    config = config or TextConfig()
    normalized = unicodedata.normalize(config.unicode_form, str(text))
    if config.remove_tatweel:
        normalized = normalized.replace(TATWEEL, "")
    if config.remove_diacritics:
        normalized = DIACRITIC_PATTERN.sub("", normalized)
    if not config.keep_punctuation:
        normalized = PUNCTUATION_PATTERN.sub("", normalized)
    if config.collapse_spaces:
        normalized = SPACE_PATTERN.sub(" ", normalized)
    if config.strip_text:
        normalized = normalized.strip()
    return normalized


def load_image_from_bytes(image_bytes: bytes) -> Image.Image:
    with BytesIO(image_bytes) as handle:
        image = Image.open(handle)
        image.load()
    return image


def resize_to_height(image: Image.Image, target_height: int) -> Image.Image:
    if image.height == target_height:
        return image
    scale = target_height / image.height
    target_width = max(1, round(image.width * scale))
    return image.resize((target_width, target_height), Image.BILINEAR)


def pad_to_width(image: Image.Image, target_width: int, fill: int = 255) -> Image.Image:
    if image.width >= target_width:
        return image
    canvas = Image.new(image.mode, (target_width, image.height), fill)
    canvas.paste(image, (0, 0))
    return canvas


def preprocess_image(image: Image.Image, config: ImageConfig) -> np.ndarray:
    if config.grayscale:
        image = image.convert("L")
    if config.target_height and config.preserve_aspect_ratio:
        image = resize_to_height(image, config.target_height)
    elif config.target_height and config.target_width:
        image = image.resize((config.target_width, config.target_height), Image.BILINEAR)
    if config.target_width:
        image = pad_to_width(image, config.target_width)
    if config.binarize:
        image = image.point(lambda value: 255 if value >= 128 else 0)
    array = np.asarray(image, dtype=np.float32)
    return array / 255.0 if config.normalize else array


def edit_distance(source: Sequence[object], target: Sequence[object]) -> int:
    previous = list(range(len(target) + 1))
    for source_index, source_item in enumerate(source, start=1):
        current = [source_index]
        for target_index, target_item in enumerate(target, start=1):
            current.append(
                min(
                    previous[target_index] + 1,
                    current[-1] + 1,
                    previous[target_index - 1] + (source_item != target_item),
                )
            )
        previous = current
    return previous[-1]


def cer(reference: str, prediction: str) -> float:
    return edit_distance(list(reference), list(prediction)) / max(len(reference), 1)


def wer(reference: str, prediction: str) -> float:
    return edit_distance(reference.split(), prediction.split()) / max(len(reference.split()), 1)


def summarize_predictions(rows) -> dict[str, float]:
    total_chars = rows["text_reference"].map(lambda value: max(len(str(value)), 1)).sum()
    total_words = rows["text_reference"].map(lambda value: max(len(str(value).split()), 1)).sum()
    char_edits = [
        edit_distance(list(str(reference)), list(str(prediction)))
        for reference, prediction in zip(rows["text_reference"], rows["text_prediction"])
    ]
    word_edits = [
        edit_distance(str(reference).split(), str(prediction).split())
        for reference, prediction in zip(rows["text_reference"], rows["text_prediction"])
    ]
    return {
        "samples": float(len(rows)),
        "cer": float(sum(char_edits) / max(total_chars, 1)),
        "wer": float(sum(word_edits) / max(total_words, 1)),
        "exact_match_rate": float(
            (rows["text_reference"].astype(str) == rows["text_prediction"].astype(str)).mean()
        ),
    }


class Charset:
    def __init__(self, characters: Sequence[str]):
        self.blank_id = 0
        self.characters = list(characters)
        self.char_to_id = {character: index + 1 for index, character in enumerate(self.characters)}
        self.id_to_char = {index + 1: character for index, character in enumerate(self.characters)}

    @classmethod
    def from_texts(cls, texts: Iterable[str]) -> "Charset":
        return cls(sorted({character for text in texts for character in str(text)}))

    @property
    def vocab_size(self) -> int:
        return len(self.characters) + 1

    def encode(self, text: str) -> list[int]:
        unknown = sorted({character for character in text if character not in self.char_to_id})
        if unknown:
            raise ValueError(f"Characters not in training charset: {''.join(unknown)!r}")
        return [self.char_to_id[character] for character in text]

    def decode(self, token_ids: Sequence[int]) -> str:
        return "".join(self.id_to_char[token_id] for token_id in token_ids if token_id != self.blank_id)


def greedy_ctc_decode(token_ids: Sequence[int], blank_id: int = 0) -> list[int]:
    decoded: list[int] = []
    previous = None
    for token_id in token_ids:
        if token_id != blank_id and token_id != previous:
            decoded.append(int(token_id))
        previous = token_id
    return decoded


def read_manifest(path: Path) -> pd.DataFrame:
    import pandas as pd

    return pd.read_csv(path)


def manifest_overview(manifest_dir: Path) -> pd.DataFrame:
    import pandas as pd

    rows = []
    for path in sorted(manifest_dir.glob("*.csv")):
        frame = pd.read_csv(path)
        rows.append(
            {
                "file": path.name,
                "rows": len(frame),
                "dataset": frame["dataset"].iloc[0] if "dataset" in frame and len(frame) else "",
                "split": frame["split"].iloc[0] if "split" in frame and len(frame) else "",
            }
        )
    return pd.DataFrame(rows)


def sample_ahcd_zip(zip_path: Path, limit: int = 256, image_config: ImageConfig | None = None):
    image_config = image_config or ImageConfig(target_height=32, target_width=32, preserve_aspect_ratio=False)
    images: list[np.ndarray] = []
    labels: list[int] = []
    sample_ids: list[str] = []
    with zipfile.ZipFile(zip_path) as archive:
        members = [name for name in sorted(archive.namelist()) if name.endswith(".png")]
        for member in members[:limit]:
            match = AHCD_FILE_PATTERN.search(member)
            if not match:
                continue
            with archive.open(member) as handle:
                image = Image.open(BytesIO(handle.read()))
                image.load()
            images.append(preprocess_image(image, image_config)[None, :, :])
            labels.append(int(match.group("label")) - 1)
            sample_ids.append(match.group("sample_id"))
    return np.stack(images).astype(np.float32), np.asarray(labels, dtype=np.int64), sample_ids


def safe_read_csv(path: Path) -> pd.DataFrame:
    import pandas as pd

    if not path.exists():
        return pd.DataFrame()
    return pd.read_csv(path)


def save_csv_rows(path: Path, fieldnames: Sequence[str], rows: Sequence[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def output_length_after_crnn(input_width: int) -> int:
    width = math.floor(input_width / 2)
    width = math.floor(width / 2)
    return max(width - 1, 1)