Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| from pathlib import Path | |
| from model import Model, DecoderType | |
| from preprocessor import Preprocessor | |
| from dataloader_iam import Batch | |
| import tensorflow as tf | |
| def get_img_size(line_mode: bool = False) -> tuple[int, int]: | |
| """ | |
| Auxiliary method that sets the height and width. | |
| Height is fixed while width is set according to the Model used. | |
| """ | |
| if line_mode: | |
| return 256, get_img_height() | |
| return 128, get_img_height() | |
| def get_img_height() -> int: | |
| """ | |
| Auxiliary method that sets the fixed height for the Neural Network. | |
| """ | |
| return 32 | |
| def infer(line_mode: bool, model: Model, fn_img: str) -> None: | |
| """ | |
| Auxiliary method that does inference using the pretrained models: | |
| Recognizes text in an image given its path. | |
| """ | |
| img = cv2.imread(fn_img, cv2.IMREAD_GRAYSCALE) | |
| assert img is not None | |
| preprocessor = Preprocessor(get_img_size(line_mode), dynamic_width=True, padding=16) | |
| img = preprocessor.process_img(img) | |
| batch = Batch([img], None, 1) | |
| recognized, probability = model.infer_batch(batch, True) | |
| return recognized, probability | |
| def main(image_path: str, model_path: str, decoder_type: DecoderType): | |
| """ | |
| Main function to load the model, perform inference on the input image, | |
| and print the result. | |
| """ | |
| # Load the model | |
| char_list_path = model_path + "/charList.txt" | |
| model = Model(list(open(char_list_path).read()), model_path, decoder_type, must_restore=True) | |
| # Perform inference | |
| recognized, probability = infer(model_path.endswith('line-model'), model, image_path) | |
| # Print the results | |
| print("Recognized Text:", recognized[0]) | |
| print("Probability:", probability[0]) | |
| if __name__ == "__main__": | |
| # Example usage | |
| # Define the image path, model directory, and decoder type here | |
| image_path = 'word.png' # Update this path | |
| model_path = '../model/word-model' # or '../model/line-model' depending on your model | |
| decoder_type = DecoderType.BestPath # Change as needed: BestPath, BeamSearch, WordBeamSearch | |
| # Call the main function with the specified parameters | |
| main(image_path, model_path, decoder_type) | |