Blame for predict_image.py

"""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')
        # Preserve the actual six-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()