| """Pinned H5 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.4-CUDA-TF32' |
| CORE_HASHES = { |
| 'model.safetensors': 'a00ac115f20a9b00ababec445c526771e0d6dcf8ad41cad30049bcaa40fc0ec8', |
| 'config.json': '7691e69a4e4cb70416c2e12bbefea3a90671790226b2c0d22ede944c48e11f05', |
| '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') |
| |
| 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() |
|
|
