Duplication in recognition model

[kwang19113]

I train images with the same format input like that. This problem rarely occurred but with thousands of images there will be a couple of them. I've export the model to onnx and infer with the code below

import onnxruntime as ort
import numpy as np
import cv2
import string

def load_model(model_path):
    """
    Load the ONNX model from the specified path.
    """
    session = ort.InferenceSession(model_path)
    return session

def preprocess_image_ocr(image, input_shape):
    """
    Preprocess the image to match the model's input shape.
    Args:
        image_path (str): Path to the input image.
        input_shape (tuple): Expected shape of the model input (channels, height, width).
    Returns:
        np.ndarray: Preprocessed image as a NumPy array.
    """
    image = cv2.imread(image)
    #cv2.imshow("image_ori",image)
    # Resize the image to match the model's expected input dimensions
    image = cv2.resize(image, (100, 32),interpolation = cv2.INTER_CUBIC)
    #cv2.imshow("image",image)
    #cv2.waitKey(0)
    # Convert the image to RGB format if needed
    #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # Normalize the pixel values to [0, 1]
    image = image.astype(np.float32)
    # Rearrange dimensions to match model input (channels, height, width)
    image = np.transpose(image, (2, 0, 1))
    # Add a batch dimension (1, channels, height, width)
    image = np.expand_dims(image, axis=0)
    return image

def run_inference_ocr(model_path, input_data):
    """
    Run inference using the ONNX model.
    Args:
        model_path (str): Path to the ONNX model.
        input_data (np.ndarray): Preprocessed input data.
    Returns:
        np.ndarray: Model output.
    """
    # Load the model
    session = load_model(model_path)
    # Get the input name for the model
    input_name = session.get_inputs()[0].name
    # Run inference
    output = session.run(None, {input_name: input_data})
    return output

def format_output_ocr(output):
    """
    Convert the model output to a string of numbers and characters.
    Args:
        output (list): Raw model output.
    Returns:
        str: Formatted output as a string of numbers and characters.
    """
    # Find the index of the maximum value in each output vector to get the predicted character
    character_indices = np.argmax(output[0], axis=-1)
    conf_scores = np.max(output[0], axis=-1)
    # Map indices to characters (assuming the provided character mapping)
    characters = " 0123456789abcdefghijklmnopqrstuvwxyz"
    result = ''.join([characters[i] for i in character_indices.flatten() if i < len(characters)])
    result_with_scores = ""
    for i in range(len(result)):
        result_with_scores += result[i] +" score:"+str(conf_scores[0][i])+"\n"
    return result, result_with_scores

def main():
    # Path to the ONNX model
    model_path = 'ocr.onnx'
    # Path to the input image
    image_path = 'det_infer/bbbb_wrapped_0.png'
    # Define the expected input shape of the model (channels, height, width)
    input_shape = (3, 32, 100)  # Adjust this to match your model's input requirements
    
    # Preprocess the image
    
    input_data = preprocess_image_ocr(image_path, input_shape)
    
    # Run inference
    output = run_inference_ocr(model_path, input_data)
    
    # Format and print the output
    formatted_output,result_with_scores = format_output_ocr(output)
    print("Model output:\n", formatted_output)
    print(result_with_scores)

if __name__ == "__main__":
    main()

[GreatV]

please refer to https://github.com/PaddlePaddle/PaddleOCR/blob/main/tools/infer/predict_rec.py