| |
| |
| |
|
|
| import argparse |
|
|
| from nemotron_ocr.inference.pipeline import NemotronOCR |
|
|
|
|
| def main(image_path, merge_level, no_visualize, model_dir, lang): |
| if model_dir is not None: |
| ocr_pipeline = NemotronOCR(model_dir=model_dir) |
| else: |
| ocr_pipeline = NemotronOCR(lang=lang) |
|
|
| predictions = ocr_pipeline(image_path, merge_level=merge_level, visualize=not no_visualize) |
|
|
| print(f"Found {len(predictions)} text regions.") |
| for pred in predictions: |
| 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}]" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Run OCR inference and annotate 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 (word, sentence, paragraph).", |
| ) |
| parser.add_argument("--no-visualize", action="store_true", help="Do not save the annotated image.") |
| parser.add_argument( |
| "--model-dir", |
| type=str, |
| default=None, |
| help="Path to a directory with detector.pth, recognizer.pth, relational.pth, charset.txt. " |
| "If omitted, weights are downloaded from Hugging Face (default: v2 multilingual).", |
| ) |
| parser.add_argument( |
| "--lang", |
| type=str, |
| choices=["en", "multi", "v1"], |
| default=None, |
| help="Hub checkpoint when --model-dir is omitted: en=v2 English, multi=v2 multilingual (default), v1=legacy.", |
| ) |
| 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, |
| ) |
|
|
