| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
|
|
|
|
| class HTRProcessor: |
| model_input_names = ["pixel_values"] |
|
|
| @classmethod |
| def register_for_auto_class(cls, auto_class="AutoProcessor"): |
| """Compatibility with transformers AutoProcessor (no-op for custom processor).""" |
| pass |
|
|
| def __init__( |
| self, |
| characters: list[str], |
| image_height: int = 64, |
| image_max_width: int = 3072, |
| width_stride: int = 32, |
| resample: str = "bilinear", |
| ) -> None: |
| self.characters = characters |
| self.image_height = int(image_height) |
| self.image_max_width = int(image_max_width) |
| self.width_stride = int(width_stride) |
| self.resample = resample |
| self.id_to_char = {idx + 1: char for idx, char in enumerate(self.characters)} |
|
|
| @staticmethod |
| def _resolve_file( |
| path_or_repo_id: str, filename: str, local_files_only: bool |
| ) -> str: |
| candidate = Path(path_or_repo_id) / filename |
| if candidate.exists(): |
| return str(candidate) |
| return hf_hub_download( |
| repo_id=path_or_repo_id, |
| filename=filename, |
| local_files_only=local_files_only, |
| ) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| pretrained_model_name_or_path: str, |
| local_files_only: bool = False, |
| **_: dict[str, Any], |
| ) -> "HTRProcessor": |
| cfg_path = cls._resolve_file( |
| pretrained_model_name_or_path, "preprocessor_config.json", local_files_only |
| ) |
| with open(cfg_path, "r", encoding="utf-8") as f: |
| cfg = json.load(f) |
|
|
| vocab_filename = cfg.get("vocab_file", "alphabet.json") |
| vocab_path = cls._resolve_file( |
| pretrained_model_name_or_path, vocab_filename, local_files_only |
| ) |
| with open(vocab_path, "r", encoding="utf-8") as f: |
| vocab_data = json.load(f) |
|
|
| if isinstance(vocab_data, dict) and "characters" in vocab_data: |
| characters = vocab_data["characters"] |
| elif isinstance(vocab_data, list): |
| characters = vocab_data |
| else: |
| raise ValueError( |
| "Unsupported vocab file format. Expected list or {'characters': [...]} ." |
| ) |
|
|
| return cls( |
| characters=characters, |
| image_height=cfg.get("image_height", 64), |
| image_max_width=cfg.get("image_max_width", 3072), |
| width_stride=cfg.get("width_stride", 32), |
| resample=cfg.get("resample", "bilinear"), |
| ) |
|
|
| def save_pretrained(self, save_directory: str) -> None: |
| os.makedirs(save_directory, exist_ok=True) |
| vocab_path = os.path.join(save_directory, "alphabet.json") |
| with open(vocab_path, "w", encoding="utf-8") as f: |
| json.dump({"characters": self.characters}, f, ensure_ascii=False, indent=2) |
|
|
| preprocessor_cfg = { |
| "processor_class": self.__class__.__name__, |
| "vocab_file": "alphabet.json", |
| "image_height": self.image_height, |
| "image_max_width": self.image_max_width, |
| "width_stride": self.width_stride, |
| "resample": self.resample, |
| } |
| with open( |
| os.path.join(save_directory, "preprocessor_config.json"), |
| "w", |
| encoding="utf-8", |
| ) as f: |
| json.dump(preprocessor_cfg, f, ensure_ascii=False, indent=2) |
|
|
| def _load_pil(self, image: str | Image.Image | np.ndarray) -> Image.Image: |
| if isinstance(image, Image.Image): |
| return image.convert("L") |
| if isinstance(image, np.ndarray): |
| if image.ndim == 2: |
| return Image.fromarray(image).convert("L") |
| if image.ndim == 3: |
| return Image.fromarray(image).convert("L") |
| raise ValueError(f"Unsupported ndarray shape: {image.shape}") |
| if isinstance(image, str): |
| return Image.open(image).convert("L") |
| raise TypeError(f"Unsupported image input type: {type(image)}") |
|
|
| def _preprocess_image(self, image: str | Image.Image | np.ndarray) -> np.ndarray: |
| img = self._load_pil(image) |
| w, h = img.size |
| if h <= 0: |
| raise ValueError("Input image has invalid height.") |
| scale = self.image_height / float(h) |
| new_w = max(1, int(w * scale)) |
|
|
| resample_map = { |
| "nearest": Image.Resampling.NEAREST, |
| "bilinear": Image.Resampling.BILINEAR, |
| "bicubic": Image.Resampling.BICUBIC, |
| "lanczos": Image.Resampling.LANCZOS, |
| } |
| pil_resample = resample_map.get( |
| self.resample.lower(), Image.Resampling.BILINEAR |
| ) |
| img = img.resize((new_w, self.image_height), resample=pil_resample) |
| arr = np.array(img) |
|
|
| if new_w > self.image_max_width: |
| arr = arr[:, : self.image_max_width] |
| new_w = self.image_max_width |
|
|
| if new_w % self.width_stride != 0: |
| aligned_w = ((new_w // self.width_stride) + 1) * self.width_stride |
| pad_width = aligned_w - new_w |
| arr = np.pad( |
| arr, |
| ((0, 0), (0, pad_width)), |
| mode="constant", |
| constant_values=0, |
| ) |
| new_w = aligned_w |
|
|
| arr = arr.astype(np.float32) / 255.0 |
| if arr.ndim == 2: |
| arr = np.expand_dims(arr, axis=-1) |
| return arr.transpose(2, 0, 1).astype(np.float32) |
|
|
| def __call__( |
| self, |
| images: str | Image.Image | np.ndarray | list[str | Image.Image | np.ndarray], |
| return_tensors: str = "pt", |
| **_: dict[str, Any], |
| ) -> dict[str, Any]: |
| batch_images = images if isinstance(images, list) else [images] |
| pixel_values = np.stack( |
| [self._preprocess_image(img) for img in batch_images], axis=0 |
| ) |
| if return_tensors == "pt": |
| return {"pixel_values": torch.from_numpy(pixel_values)} |
| if return_tensors == "np": |
| return {"pixel_values": pixel_values} |
| raise ValueError("Supported return_tensors values are 'pt' and 'np'.") |
|
|
| @staticmethod |
| def _ctc_greedy_decode( |
| logits_tnc: np.ndarray, blank_idx: int = 0 |
| ) -> list[list[int]]: |
| preds = np.argmax(logits_tnc, axis=2) |
| _, batch_size, _ = logits_tnc.shape |
| decoded: list[list[int]] = [] |
| for n in range(batch_size): |
| seq = preds[:, n] |
| chars: list[int] = [] |
| prev = blank_idx |
| for idx in seq: |
| token = int(idx) |
| if token != blank_idx and token != prev: |
| chars.append(token) |
| prev = token |
| decoded.append(chars) |
| return decoded |
|
|
| def batch_decode( |
| self, |
| logits: torch.Tensor | np.ndarray, |
| blank_idx: int = 0, |
| logit_layout: str = "ntc", |
| ) -> list[str]: |
| logits_np = ( |
| logits.detach().cpu().numpy() |
| if isinstance(logits, torch.Tensor) |
| else logits |
| ) |
| if logits_np.ndim != 3: |
| raise ValueError(f"Expected logits rank 3, got shape {logits_np.shape}.") |
|
|
| if logit_layout == "ntc": |
| logits_tnc = np.transpose(logits_np, (1, 0, 2)) |
| elif logit_layout == "tnc": |
| logits_tnc = logits_np |
| else: |
| raise ValueError("logit_layout must be 'ntc' or 'tnc'.") |
|
|
| decoded_ids = self._ctc_greedy_decode(logits_tnc, blank_idx=blank_idx) |
| return [ |
| "".join( |
| self.id_to_char.get(token, "") |
| for token in seq |
| if token in self.id_to_char |
| ) |
| for seq in decoded_ids |
| ] |
|
|
