File size: 2,804 Bytes
4d53871 | 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 | #!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import argparse
from nemotron_ocr.inference.pipeline_v2 import NemotronOCRV2
def main(image_path, merge_level, no_visualize, model_dir, lang,
detector_only, skip_relational):
kwargs = {}
if model_dir is not None:
kwargs["model_dir"] = model_dir
else:
kwargs["lang"] = lang
if detector_only:
kwargs["detector_only"] = True
if skip_relational:
kwargs["skip_relational"] = True
ocr = NemotronOCRV2(**kwargs)
predictions = ocr(image_path, merge_level=merge_level)
print(f"Found {len(predictions)} text regions.")
for pred in predictions:
if "text" in pred:
print(
f" - Text: '{pred['text']}', "
f"Confidence: {pred['confidence']:.2f}, "
f"Bbox: [left={pred['left']:.4f}, upper={pred['upper']:.4f}, "
f"right={pred['right']:.4f}, lower={pred['lower']:.4f}]"
)
else:
print(
f" - Confidence: {pred['confidence']:.2f}, "
f"Bbox: [left={pred['left']:.4f}, upper={pred['upper']:.4f}, "
f"right={pred['right']:.4f}, lower={pred['lower']:.4f}]"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run OCR inference on an image.")
parser.add_argument("image_path", type=str, help="Path to the input image.")
parser.add_argument(
"--merge-level",
type=str,
choices=["word", "sentence", "paragraph"],
default="paragraph",
help="Merge level for OCR output (default: paragraph).",
)
parser.add_argument("--no-visualize", action="store_true", help="(unused, kept for compat)")
parser.add_argument(
"--model-dir", type=str, default=None,
help="Local checkpoint directory. If omitted, downloads from Hugging Face.",
)
parser.add_argument(
"--lang", type=str, choices=["en", "multi", "v1"], default=None,
help="Hub checkpoint: en, multi (default), or v1.",
)
parser.add_argument(
"--detector-only", action="store_true",
help="Run detector only — returns boxes without text.",
)
parser.add_argument(
"--skip-relational", action="store_true",
help="Skip relational model — returns per-word text without reading order.",
)
args = parser.parse_args()
main(
args.image_path,
merge_level=args.merge_level,
no_visualize=args.no_visualize,
model_dir=args.model_dir,
lang=args.lang,
detector_only=args.detector_only,
skip_relational=args.skip_relational,
)
|