""" captCHAD Inference Engine Supports PyTorch, Safetensors, and ONNX Runtime across FP32, FP16, INT8, FP8, and INT4 quantizations. Usage: python inference.py sample.png python inference.py sample.png --engine onnx --quant int8 python inference.py sample.png --engine onnx --quant fp16 python inference.py sample.png --engine safetensors --quant fp8 python inference.py sample.png --engine safetensors --quant int4 """ import os import sys import argparse import time import numpy as np from PIL import Image CHARSET = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" IDX2CHAR = {i + 1: ch for i, ch in enumerate(CHARSET)} BLANK_IDX = 0 def preprocess_image(image_path: str) -> np.ndarray: """Preprocess image to normalized float32 array (1, 3, 64, 192).""" with Image.open(image_path) as img: if img.mode == "RGBA": bg = Image.new("RGB", img.size, (255, 255, 255)) bg.paste(img, mask=img.split()[3]) img = bg elif img.mode != "RGB": img = img.convert("RGB") img = img.resize((192, 64), Image.BILINEAR) arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) # (3, 64, 192) # Normalize to [-1.0, 1.0] arr = (arr - 127.5) / 127.5 return arr[np.newaxis, :, :, :] # (1, 3, 64, 192) def ctc_decode_greedy(tokens: list[int]) -> str: """Standard CTC collapse: drop consecutive duplicates and blank tokens.""" res = [] prev = None for t in tokens: if t != prev and t != BLANK_IDX: if t in IDX2CHAR: res.append(IDX2CHAR[t]) prev = t return "".join(res) class captCHADPredictor: def __init__(self, engine: str = "onnx", quant: str = "fp32", weights_path: str = None): self.engine = engine.lower() self.quant = quant.lower() dir_path = os.path.dirname(os.path.abspath(__file__)) # Resolve weights path if not given if weights_path is None: if self.engine == "onnx": if self.quant == "int8": weights_path = os.path.join(dir_path, "captchad_int8.onnx") elif self.quant == "fp16": weights_path = os.path.join(dir_path, "captchad_fp16.onnx") else: weights_path = os.path.join(dir_path, "captchad.onnx") elif self.engine == "safetensors": if self.quant == "fp16": weights_path = os.path.join(dir_path, "model_fp16.safetensors") elif self.quant == "fp8": weights_path = os.path.join(dir_path, "model_fp8.safetensors") elif self.quant == "int4": weights_path = os.path.join(dir_path, "model_int4.safetensors") else: weights_path = os.path.join(dir_path, "model.safetensors") else: # pytorch if self.quant == "fp16": weights_path = os.path.join(dir_path, "captchad_fp16.pt") elif self.quant == "int8": weights_path = os.path.join(dir_path, "captchad_int8.pt") elif self.quant == "fp8": weights_path = os.path.join(dir_path, "captchad_fp8.pt") elif self.quant == "int4": weights_path = os.path.join(dir_path, "captchad_int4.pt") else: weights_path = os.path.join(dir_path, "captchad.pt") self.weights_path = weights_path if self.engine == "onnx": import onnxruntime as ort opts = ort.SessionOptions() opts.intra_op_num_threads = min(4, os.cpu_count() or 4) self.session = ort.InferenceSession(weights_path, sess_options=opts) self.input_name = self.session.get_inputs()[0].name elif self.engine in ("pytorch", "safetensors"): import torch from model import captCHAD, decode_beam_search_single self.torch = torch self.decode_beam = decode_beam_search_single self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model = captCHAD(num_classes=len(CHARSET) + 1) if weights_path.endswith("int8.pt"): ckpt = torch.load(weights_path, map_location=self.device, weights_only=False) self.model = ckpt["model"] if "model" in ckpt else ckpt elif weights_path.endswith(".safetensors"): from safetensors.torch import load_file state = load_file(weights_path) if "int4" in weights_path: # Dequantize packed int4 weights base_ckpt = torch.load(os.path.join(dir_path, "captchad.pt"), map_location="cpu") orig_shapes = {k: v.shape for k, v in base_ckpt["model_state_dict"].items()} restored = {} for k, v in state.items(): if k.endswith(".packed_int4"): bname = k[:-len(".packed_int4")] sc = state[f"{bname}.scale"].squeeze() oshape = orig_shapes[bname] low = (v & 0x0F).to(torch.int8) - 7 high = ((v >> 4) & 0x0F).to(torch.int8) - 7 unpacked = torch.empty(len(v) * 2, dtype=torch.int8) unpacked[0::2], unpacked[1::2] = low, high nel = 1 for d in oshape: nel *= d restored[bname] = (unpacked[:nel].to(torch.float32) * sc).reshape(oshape) elif k.endswith(".scale"): continue else: restored[k] = v.float() if v.is_floating_point() else v state = restored else: state = {k: v.to(torch.float32) if v.is_floating_point() else v for k, v in state.items()} self.model.load_state_dict(state) else: ckpt = torch.load(weights_path, map_location=self.device) state = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt state = {k: v.to(torch.float32) if v.is_floating_point() else v for k, v in state.items()} self.model.load_state_dict(state) self.model.to(self.device) self.model.eval() def predict(self, image_path: str, use_beam: bool = False) -> tuple[str, float]: t0 = time.perf_counter() inp = preprocess_image(image_path) if self.engine == "onnx": logits = self.session.run(None, {self.input_name: inp})[0] preds = np.argmax(logits[:, 0, :], axis=-1).tolist() text = ctc_decode_greedy(preds) else: t = self.torch.from_numpy(inp).to(self.device) with self.torch.no_grad(): logits = self.model(t) # (48, 1, 63) if use_beam: log_probs = logits[:, 0, :].log_softmax(dim=-1) beam_res = self.decode_beam(log_probs, beam_width=15) text = beam_res[0][0] if beam_res else "" else: preds = logits.argmax(dim=-1)[:, 0].tolist() text = ctc_decode_greedy(preds) latency_ms = (time.perf_counter() - t0) * 1000 return text, latency_ms def main(): parser = argparse.ArgumentParser(description="captCHAD Multi-Format Inference Engine") parser.add_argument("image", nargs="?", default="sample.png", help="Path to input image") parser.add_argument("--engine", choices=["onnx", "pytorch", "safetensors"], default="onnx", help="Inference engine") parser.add_argument("--quant", choices=["fp32", "fp16", "int8", "fp8", "int4"], default="fp32", help="Precision format") parser.add_argument("--weights", type=str, default=None, help="Custom weights file path") parser.add_argument("--beam", action="store_true", help="Use CTC beam search (PyTorch only)") args = parser.parse_args() if not os.path.exists(args.image): print(f"Error: image not found at '{args.image}'.") sys.exit(1) predictor = captCHADPredictor(engine=args.engine, quant=args.quant, weights_path=args.weights) pred_text, latency = predictor.predict(args.image, use_beam=args.beam) print(f"Image: {args.image}") print(f"Engine: {args.engine.upper()} ({args.quant.upper()})") print(f"Weights: {os.path.basename(predictor.weights_path)}") print(f"Prediction: {pred_text}") print(f"Latency: {latency:.2f} ms") if __name__ == "__main__": main()