RenAI / main.py
Arsh124's picture
Initial RenAI app
ebcc7d1
from skimage.io import imread, imsave
from skimage.color import rgb2gray
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from skimage.transform import resize
from utils.preprocessing import preprocessImage, postProcessImage, process_segment_and_crop_image
from utils.line_segmentation import segment_image_to_lines
from configs import unet_enabled
from utils.helper import load_images_from_json
from inference import Inference
from configs import model_path, processor_path, unet_model_path
from utils.postprocessing import PostProcessing
from loguru import logger
def RenAITranscription(image, post_processing_enabled=False,unet_enabled=False):
# 1- preprocessing
org_img = imread(image)[: , : ,:]
logger.info(f'Image Dimensions : {org_img.shape[0]} x {org_img.shape[1]}')
intial_process_image = preprocessImage(org_img)
if unet_enabled:
logger.info("Masked based segmentation and cropping enabled...")
cropped_img = process_segment_and_crop_image(unet_model_path, org_img, intial_process_image, padding=10, min_contour_area=100)
processed_image = postProcessImage(cropped_img)
logger.info(f"Image cropped and Pre-processed successfully.....")
else:
logger.info("Image Preprocessing started......")
processed_image = postProcessImage(intial_process_image)
logger.info(f"Image Pre-processed successfully.....")
# 2 - Line segmentation Algorithm
line_segments = segment_image_to_lines(processed_image, base_key="line",ct=0)
# 3 - Model Inference
transciption_generator = Inference(
model_path=model_path,
processor_path=processor_path,
target_size=(256, 64),
batch_size=32
)
result = transciption_generator.generate_texts_from_images(line_segments)
# Generated texts
for key, value in result.items():
print(f"{key}: {value['transcription']}")
# 4 - Post processing
# Dictionary based fuzzy matching
if post_processing_enabled:
for key, value in result.items():
corrected = PostProcessing(value['transcription'])
result[key]['post_processed'] = corrected
print(f"{key}: {value['post_processed']}")
print(result)
logger.info("Transcription completed successfully!")
return result
if __name__ == "__main__":
RenAITranscription("1.png", post_processing_enabled=False, unet_enabled=False)