File size: 6,689 Bytes
6536f3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """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()
|