|
|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
|
|
|
|
from nemo_retriever_ocr.inference.pipeline import NemoRetrieverOCR |
|
|
|
|
|
|
|
|
def main(image_path, merge_level, no_visualize, model_dir): |
|
|
ocr_pipeline = NemoRetrieverOCR() |
|
|
|
|
|
predictions = ocr_pipeline(image_path, merge_level=merge_level, visualize=not no_visualize) |
|
|
|
|
|
print(f"Found {len(predictions)} text regions.") |
|
|
|
|
|
|
|
|
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, |
|
|
help="Path to the model checkpoints.", |
|
|
default="./checkpoints", |
|
|
) |
|
|
args = parser.parse_args() |
|
|
|
|
|
main( |
|
|
args.image_path, |
|
|
merge_level=args.merge_level, |
|
|
no_visualize=args.no_visualize, |
|
|
model_dir=args.model_dir, |
|
|
) |
|
|
|
