"""Pinned owner-selected H5.1 raster inference example; native model outputs map to AX PDF SDK IDs.""" from __future__ import annotations import argparse import hashlib import importlib.metadata import json import math from pathlib import Path import torch from PIL import Image from transformers import RTDetrImageProcessor, RTDetrV2ForObjectDetection REPO_ID = 'AutomatosX/AXDocLayout-M-v1.5-CUDA-TF32' CORE_HASHES = {'config.json': '7691e69a4e4cb70416c2e12bbefea3a90671790226b2c0d22ede944c48e11f05', 'model.safetensors': 'd0949704220ee50d94359cb6417e3481ac4ef5078a9f63795989e9f5d76e9b7e', 'preprocessor_config.json': 'cd38cd59999e7a95d68e487fbe5132df3d4e5c32a0836add57e6126ba0c4eaf1', 'sdk-taxonomy.json': '81e7908ee7c8bbb6e00905a98ed43c988128faa38c0b1d55d20948f78e40e351'} def digest(path: Path) -> str: with path.open('rb') as stream: return hashlib.file_digest(stream, 'sha256').hexdigest() def validate_threshold(value: float) -> float: if isinstance(value, bool) or not isinstance(value, (float, int)): raise TypeError('score_threshold_invalid') if not math.isfinite(value) or not 0 <= value <= 1: raise ValueError('score_threshold_invalid') return float(value) def load_bundle(*, model_dir: Path, device: str = 'cuda') -> tuple: """Load exact local files with an explicit device; never fall back from CUDA.""" if device not in ('cuda', 'cpu'): raise ValueError('device_invalid') if importlib.metadata.version('transformers') != '4.53.3': raise RuntimeError('transformers_version_mismatch') if torch.__version__.split('+')[0] != '2.12.1': raise RuntimeError('torch_version_mismatch') for name, expected in CORE_HASHES.items(): path = model_dir / name if not path.is_file() or digest(path) != expected: raise RuntimeError('model_bundle_hash_mismatch') if device == 'cuda': if not torch.cuda.is_available(): raise RuntimeError('cuda_required_no_cpu_fallback') if torch.cuda.get_device_capability()[0] < 8: raise RuntimeError('ampere_or_newer_required_for_tf32_profile') # Preserve the actual three-panel CUDA evaluation configuration. torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = True model, loading = RTDetrV2ForObjectDetection.from_pretrained( model_dir, local_files_only=True, output_loading_info=True, ) if any(loading[key] for key in ('missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs')): raise RuntimeError('model_state_mismatch') binding = json.loads((model_dir / 'sdk-taxonomy.json').read_text()) if binding['taxonomy'] != 'ax-document-unified-25/1' or binding['model_sdk_class_ids'] != list(range(1, 18)): raise RuntimeError('sdk_binding_mismatch') definitions = {row['id']: row for row in binding['classes']} for native, sdk in binding['native_index_to_sdk_id'].items(): if model.config.id2label[int(native)] != definitions[sdk]['name']: raise RuntimeError('sdk_name_mismatch') processor = RTDetrImageProcessor.from_pretrained(model_dir, local_files_only=True) return model.float().to(device).eval(), processor, binding def predict_rgb(*, rgb: Image.Image, model: object, processor: object, binding: dict, threshold: float = .30) -> list[dict]: """Return original-raster xyxy boxes with SDK IDs1-17 and finite scores.""" threshold = validate_threshold(threshold) if not isinstance(rgb, Image.Image) or rgb.mode != 'RGB' or min(rgb.size) < 1: raise ValueError('decoded_rgb_image_required') device = next(model.parameters()).device inputs = processor(images=rgb, return_tensors='pt')['pixel_values'] if tuple(inputs.shape) != (1, 3, 640, 640): raise RuntimeError('preprocessor_shape_mismatch') with torch.inference_mode(): outputs = model(pixel_values=inputs.to(device)) predictions = processor.post_process_object_detection( outputs, threshold=0., target_sizes=torch.tensor([[rgb.height, rgb.width]], device=device), )[0] definitions = {row['id']: row['name'] for row in binding['classes']} result = [] for native, score, box in zip(predictions['labels'].tolist(), predictions['scores'].tolist(), predictions['boxes'].tolist(), strict=True): if not all(math.isfinite(v) for v in [score, *box]) or not 0 <= score <= 1: raise RuntimeError('nonfinite_or_invalid_prediction') if score < threshold: continue sdk = binding['native_index_to_sdk_id'][str(native)] clipped = [min(max(box[0], 0.), rgb.width), min(max(box[1], 0.), rgb.height), min(max(box[2], 0.), rgb.width), min(max(box[3], 0.), rgb.height)] if clipped[2] <= clipped[0] or clipped[3] <= clipped[1]: continue result.append({'class_id': sdk, 'class_name': definitions[sdk], 'score': score, 'bbox_xyxy': clipped}) return result def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--model-dir', type=Path, required=True) parser.add_argument('--image', type=Path, required=True) parser.add_argument('--output', type=Path, required=True) parser.add_argument('--device', choices=('cuda', 'cpu'), default='cuda') parser.add_argument('--threshold', type=float, default=.30) args = parser.parse_args() validate_threshold(args.threshold) if args.output.exists(): raise RuntimeError('output_already_exists') try: with Image.open(args.image) as image: if image.format not in ('PNG', 'JPEG', 'TIFF', 'WEBP', 'BMP'): raise ValueError('raster_format_required') rgb = image.convert('RGB') except (OSError, ValueError): raise RuntimeError('raster_image_read_failed') from None model, processor, binding = load_bundle(model_dir=args.model_dir, device=args.device) detections = predict_rgb(rgb=rgb, model=model, processor=processor, binding=binding, threshold=args.threshold) result = {'model_id': REPO_ID, 'model_sha256': CORE_HASHES['model.safetensors'], 'taxonomy': binding['taxonomy'], 'width': rgb.width, 'height': rgb.height, 'coordinate_space': 'original-raster-pixel-xyxy', 'device': args.device, 'score_threshold': args.threshold, 'detections': detections} args.output.write_text(json.dumps(result, sort_keys=True, allow_nan=False) + '\n') if __name__ == '__main__': main()