app.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import cv2
import numpy as np
from PIL import Image, ImageOps
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow.data as tfd
import io
import csv
import os
import tempfile
import subprocess
import json
from typing import List, Dict, Any
import logging
import pandas as pd

# Load back
with open("vocab.json", "r") as f:
    vocab_loaded = json.load(f)

# Recreate StringLookup
char_to_num = layers.StringLookup(vocabulary=vocab_loaded, mask_token=None)
num_to_char = layers.StringLookup(vocabulary=char_to_num.get_vocabulary(),
                                  mask_token=None, invert=True)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Spanish OCR API", description="API for Spanish text recognition from images")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Global variables for model and configurations
inference_model = None
# char_to_num = None
# num_to_char = None
n_classes=50
IMG_WIDTH = 200
IMG_HEIGHT = 50
MAX_LABEL_LENGTH = None
# AUTOTUNE
AUTOTUNE = tfd.AUTOTUNE
# Batch Size
BATCH_SIZE = 16

# Character mapping for corrections
# CHARACTER_MAPPING = {
#     'в': 'o',
#     'д': 'ñ',
#     'б': 'i',
#     'В': 'e',
#     'а': 'a'
# }

class CTCLayer(layers.Layer):
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.loss_fn = keras.backend.ctc_batch_cost
    def call(self, y_true, y_pred):
        batch_len = tf.cast(tf.shape(y_true)[0], dtype='int64')
        input_len = tf.cast(tf.shape(y_pred)[1], dtype='int64') * tf.ones(shape=(batch_len, 1), dtype='int64')
        label_len = tf.cast(tf.shape(y_true)[1], dtype='int64') * tf.ones(shape=(batch_len, 1), dtype='int64')
        loss = self.loss_fn(y_true, y_pred, input_len, label_len)
        self.add_loss(loss)
        return y_pred


def load_model_and_setup():
    """Load the trained OCR model and setup character mappings"""
    global inference_model, char_to_num, num_to_char, MAX_LABEL_LENGTH
    
    try:
        # Define unique characters (update this with your actual character set)
        # unique_chars = {'e', 'j', 'Q', 'z', 'v', 'A', 'L', 't', 'V', 'O', 'c', 'q', 'l', 'a', 'ñ', 'B', 'P', ',', 'H', 'C', 'M', 'G', 's', 'r', 'T', 'd', 'g', 'p', 'D', 'S', 'N', 'b', 'm', 'u', 'o', 'f', 'I', 'x', 'R', 'y', 'n', 'i', '-', 'F', 'E', 'h'}
        
        # # Character to numeric mapping
        # char_to_num = layers.StringLookup(
        #     vocabulary=list(unique_chars),
        #     mask_token=None
        # )
        
        # # Reverse mapping
        # num_to_char = layers.StringLookup(
        #     vocabulary=char_to_num.get_vocabulary(),
        #     mask_token=None,
        #     invert=True
        # )
        
        # Load your trained model
        model_path = 'ocr_model_NEW.h5'  # Update with your model path
        if os.path.exists(model_path):
            full_model = keras.models.load_model(model_path, compile=False, custom_objects={'CTCLayer': CTCLayer})
            # Create inference model
            inference_model = keras.Model(
                inputs=full_model.get_layer(name="image").input,
                outputs=full_model.get_layer(name='dense_1').output
            )
            logger.info("Model loaded successfully")
        else:
            logger.error(f"Model file not found: {model_path}")
            raise FileNotFoundError(f"Model file not found: {model_path}")
        
        # Set MAX_LABEL_LENGTH (update with your actual value)
        MAX_LABEL_LENGTH = 24  # Update this based on your training data
        
    except Exception as e:
        logger.error(f"Error loading model: {e}")
        raise

def load_image(image_path : str):
    '''
    This function loads and preprocesses images. It first receives the image path, which is used to
    decode the image as a JPEG using TensorFlow. Then, it converts the image to a tensor and applies
    two processing functions: resizing and normalization. The processed image is then returned by
    the function.

    Argument :
        image_path : The path of the image file to be loaded.

    Return:
        image : The loaded image as a tensor.
    '''

    # Read the Image
    image = tf.io.read_file(image_path)

    # Decode the image
    decoded_image = tf.image.decode_jpeg(contents = image, channels = 1)

    # Convert image data type.
    cnvt_image = tf.image.convert_image_dtype(image = decoded_image, dtype = tf.float32)

    # Resize the image
    resized_image = tf.image.resize(images = cnvt_image, size = (IMG_HEIGHT, IMG_WIDTH))

    # Transpose
    image = tf.transpose(resized_image, perm = [1, 0, 2])

    # Convert image to a tensor.
    image = tf.cast(image, dtype = tf.float32)

    # Return loaded image
    return image

def apply_craft_detection(image_path: str, output_dir: str) -> str:
    """Apply CRAFT model for text detection"""
    try:
        # Create output directory if it doesn't exist
        os.makedirs(output_dir, exist_ok=True)
        
        # Command to run CRAFT model
        craft_command = [
            'python3', 'CRAFT_Model/CRAFT/BoundBoxFunc/test.py',
            '--cuda', '0',  # Use CPU, change to '1' if GPU available
            '--result_folder', output_dir,
            '--test_folder', os.path.dirname(image_path),
            '--trained_model', 'CRAFT_Model/CRAFT/BoundBoxFunc/weights/craft_mlt_25k.pth'
        ]
        
        # Run CRAFT detection
        result = subprocess.run(craft_command, capture_output=True, text=True)
        
        if result.returncode != 0:
            logger.error(f"CRAFT detection failed: {result.stderr}")
            raise Exception(f"CRAFT detection failed: {result.stderr}")
        
        logger.info("CRAFT detection completed successfully")
        return output_dir
        
    except Exception as e:
        logger.error(f"Error in CRAFT detection: {e}")
        raise

# def sort_bounding_boxes(bounding_box_file: str) -> List[List[int]]:
#     """Sort bounding boxes based on Spanish reading order (top to bottom, left to right)"""
#     try:
#         bounding_boxes = []
        
#         with open(bounding_box_file, 'r') as f:
#             for line in f:
#                 coords = list(map(int, line.strip().split(',')[:8]))  # Take first 8 coordinates
#                 bounding_boxes.append(coords)
        
#         # Sort by y-coordinate (top to bottom), then by x-coordinate (left to right)
#         bounding_boxes.sort(key=lambda box: (box[1], box[0]))
        
#         return bounding_boxes
        
#     except Exception as e:
#         logger.error(f"Error sorting bounding boxes: {e}")
#         return []

def count_files_in_folder(folder_path, extensions_list):
    # Initialize counter for files
    file_count = 0

    # Iterate through all files in the folder
    for filename in os.listdir(folder_path):
        # Check if the file ends with the given file extension
        for extension in extensions_list:
            if filename.lower().endswith(extension):
                file_count += 1

    return file_count

def process_bounding_boxes(file_path):
    with open(file_path, "r") as file:
        lines = file.readlines()

    # Parse bounding box coordinates
    bounding_boxes = []
    for line in lines:
        coords = list(map(int, line.strip().split(',')))
        bounding_boxes.append(coords)

    # Sort bounding boxes based on y_min value
    bounding_boxes.sort(key=lambda box: box[1])

    vertical_distance_between_lines = 10   #Change it according to the dataset, you are using
    # Group bounding boxes based on difference between max and min y_min values
    grouped_boxes = []
    current_group = []
    for box in bounding_boxes:
        if not current_group:
            current_group.append(box)
        else:
            min_y = min(current_group, key=lambda x: x[1])[1]
            max_y = max(current_group, key=lambda x: x[1])[1]
            if box[1] - min_y <= vertical_distance_between_lines:
                current_group.append(box)
            else:
                grouped_boxes.append(current_group)
                current_group = [box]

    # Append the last group
    if current_group:
        grouped_boxes.append(current_group)

    # Sort each group based on x_min value
    for group in grouped_boxes:
        group.sort(key=lambda box: box[0])

    return grouped_boxes

def sort_bounding_boxes(bounding_box_file):
    sorted_bounding_boxes = process_bounding_boxes(bounding_box_file)

    # Write sorted bounding boxes to text file in output directory
    output_file_path = f"{os.path.splitext(bounding_box_file)[0]}_sorted.txt"
    with open(output_file_path, "w") as outfile:
        for group in sorted_bounding_boxes:
            for box in group:
                outfile.write(','.join(map(str, box)) + '\n')
            outfile.write((';'))
    return output_file_path

def extract_bounding_boxes(image_path, bounding_boxes_file, output_folder, word):
    # Read the main image
    main_image = cv2.imread(image_path)
    # Create the output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Read bounding box coordinates from the text file
    with open(bounding_boxes_file, 'r') as f:
        bounding_boxes_data = f.read().split(';')
    bounding_boxes_data = bounding_boxes_data[1:]
    line=0
    for indx in range(len(bounding_boxes_data)-1):
        bounding_box_coords = bounding_boxes_data[indx].strip().split('\n')
        for cnt in range(len(bounding_box_coords)):
            coordinates_list = [int(coord) for coord in bounding_box_coords[cnt].split(',')]
            x_min, y_min, x_max, y_min, x_max, y_max, x_min, y_max = coordinates_list

            # Extract the bounding box from the main image
            bounding_box = main_image[y_min:y_max, x_min:x_max]

            # Save the bounding box as a separate image
            output_path = os.path.join(output_folder, f'{word};{line}.png')
            cv2.imwrite(output_path, bounding_box)

            word += 1
        line+=1

    return word

def pad_and_resize_images(folder_path):
    # Ensure the folder exists
    if not os.path.exists(folder_path):
        raise ValueError(f"The folder {folder_path} does not exist")

    # Define the target aspect ratio and size
    target_aspect_ratio = 4  # 1:4 aspect ratio
    target_width = 200
    target_height = 40

    # Iterate through all files in the folder
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        if os.path.isfile(file_path):
            try:
                # Open the image
                with Image.open(file_path) as img:
                    img = img.convert('L')
                    width, height = img.size
                    aspect_ratio = width / height

                    if aspect_ratio < target_aspect_ratio:
                        # Calculate padding to make aspect ratio 1:4
                        new_width = height * 4
                        padding = (new_width - width) // 2
                        padded_img = ImageOps.expand(img, border=(padding, 0, padding, 0), fill='white')
                    else:
                        padded_img = img

                    # Resize the image to 200x40
                    resized_img = padded_img.resize((target_width, target_height))

                    # Save the processed image back to the original path
                    resized_img.save(file_path)

                    print(f"Processed and replaced: {file_path}")
            except Exception as e:
                print(f"Error processing {file_path}: {e}")

def create_csv_from_folder(folder_path, csv_file_path):
    # Get a list of all files in the folder
    files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]

    # Create or overwrite the CSV file
    with open(csv_file_path, 'w', newline='') as csv_file:
        # Create a CSV writer object
        csv_writer = csv.writer(csv_file)

        # Write the header row
        csv_writer.writerow(['FILENAME', 'IDENTITY'])

        # Write data rows, excluding files with the name ".png"
        for file_name in files:
            if file_name.lower() == ".png":
                continue  # Skip files with the name ".png"

            # file_path = os.path.join(folder_path, file_name)

            # Remove the file extension (assuming it's three characters long, like '.png')
            file_name_without_extension = os.path.splitext(file_name)[0]

            csv_writer.writerow([file_name, file_name_without_extension])

    print(f'CSV file "{csv_file_path}" created successfully.')



def encode_single_sample(image_path : str, label : str):

    '''
    The function takes an image path and label as input and returns a dictionary containing the processed image tensor and the label tensor.
    First, it loads the image using the load_image function, which decodes and resizes the image to a specific size. Then it converts the given
    label string into a sequence of Unicode characters using the unicode_split function. Next, it uses the char_to_num layer to convert each
    character in the label to a numerical representation. It pads the numerical representation with a special class (n_classes)
    to ensure that all labels have the same length (MAX_LABEL_LENGTH). Finally, it returns a dictionary containing the processed image tensor
    and the label tensor.

    Arguments :
        image_path : The location of the image file.
        label      : The text to present in the image.

    Returns:
        dict : A dictionary containing the processed image and label.
    '''

    # Get the image
    image = load_image(image_path)

    # Convert the label into characters
    chars = tf.strings.unicode_split(label, input_encoding='UTF-8')

    # Convert the characters into vectors
    vecs = char_to_num(chars)

    # Pad label
    pad_size = MAX_LABEL_LENGTH - tf.shape(vecs)[0]
    vecs = tf.pad(vecs, paddings = [[0, pad_size]], constant_values=n_classes+1)

    return {'image':image, 'label':vecs}

def extract_word_images(image_path: str, bounding_boxes: List[List[int]]) -> List[np.ndarray]:
    """Extract word images from bounding boxes"""
    try:
        # Load the image
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not load image: {image_path}")
        
        word_images = []
        
        for bbox in bounding_boxes:
            # Extract coordinates
            x_coords = bbox[::2]  # Every other element starting from 0
            y_coords = bbox[1::2]  # Every other element starting from 1
            
            # Get bounding rectangle
            x_min, x_max = min(x_coords), max(x_coords)
            y_min, y_max = min(y_coords), max(y_coords)
            
            # Extract word image
            word_img = image[y_min:y_max, x_min:x_max]
            import matplotlib.pyplot as plt
            plt.plot(word_img)
            plt.savefig("image0.png")
            if word_img.size > 0:
                word_images.append(word_img)
            
        return word_images
        
    except Exception as e:
        logger.error(f"Error extracting word images: {e}")
        return []


def preprocess_word_image(word_image: np.ndarray) -> np.ndarray:
    """
    Preprocess a word image for model input.

    Steps:
    - Convert to grayscale
    - Pad to maintain aspect ratio (target 1:4)
    - Resize to (200, 50)
    - Normalize pixel values to [0, 1]
    - Transpose shape to match (height, width, channel)

    Args:
        word_image (np.ndarray): Input image as a NumPy array.

    Returns:
        np.ndarray: Preprocessed image ready for model input.
    """
    try:
        target_aspect_ratio = 4  # 1:4
        target_width = IMG_WIDTH
        target_height = IMG_HEIGHT

        # Convert color image to grayscale if needed
        if len(word_image.shape) == 3:
            word_image = cv2.cvtColor(word_image, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(word_image).convert('L')
        import matplotlib.pyplot as plt
        plt.plot(img)
        plt.savefig("image1.png")
        # Get current size and aspect ratio
        width, height = img.size
        aspect_ratio = width / height

        # Pad if aspect ratio is less than target
        if aspect_ratio < target_aspect_ratio:
            new_width = height * target_aspect_ratio
            padding = (int((new_width - width) // 2), 0)
            img = ImageOps.expand(img, border=(padding[0], 0, padding[0], 0), fill='white')

        # Resize
        img = img.resize((target_width, target_height))

        # Convert to NumPy
        img_array = np.array(img).astype(np.float32)  # shape (H, W)
        import matplotlib.pyplot as plt
        plt.plot(img)
        plt.savefig("image2.png")

        # Add channel dimension and transpose to (H, W, 1)
        img_array = img_array[:, :, np.newaxis]
        img_array = np.transpose(img_array, (1, 0, 2))  # shape (W, H, 1) → (H, W, 1)

        return img_array

    except Exception as e:
        logger.error(f"Error preprocessing word image: {e}")
        return None

def decode_prediction(pred_label: np.ndarray) -> List[str]:
    """Decode model predictions to text"""
    try:
        # Input length
        input_len = np.ones(shape=pred_label.shape[0]) * pred_label.shape[1]
        
        # CTC decode
        decode = keras.backend.ctc_decode(
            pred_label, 
            input_length=input_len, 
            greedy=True, 
        )[0][0][:, :MAX_LABEL_LENGTH]
        
        # Convert back to characters
        chars = num_to_char(decode)
        
        # Join characters
        texts = [tf.strings.reduce_join(inputs=char).numpy().decode('UTF-8') for char in chars]
        
        # Clean up text
        filtered_texts = [text.replace('[UNK]', " ").strip() for text in texts]
        
        return filtered_texts
        
    except Exception as e:
        logger.error(f"Error decoding predictions: {e}")
        return []

# Set the new size in pixels (width, height) according to your choice
def resize_images_in_folder(input_folder, new_size=(200,50)):
    # Loop through all files in the input folder
    for filename in os.listdir(input_folder):
        # Open the image
        with Image.open(os.path.join(input_folder, filename)) as img:
            # Resize the image
            resized_img = img.resize(new_size)
            # Save the resized image to the output folder
            output_filename = os.path.splitext(filename)[0] + '.png'  # Ensure output format is PNG
            resized_img.save(os.path.join(input_folder, output_filename))

def decode_pred(pred_label):

    '''
    The decode_pred function is used to decode the predicted labels generated by the OCR model.
    It takes a matrix of predicted labels as input, where each time step represents the probability
    for each character. The function uses CTC decoding to decode the numeric labels back into their
    character values. The function also removes any unknown tokens and returns the decoded texts as a
    list of strings. The function utilizes the num_to_char function to map numeric values back to their
    corresponding characters. Overall, the function is an essential step in the OCR process, as it allows
    us to obtain the final text output from the model's predictions.

    Argument :
        pred_label : These are the model predictions which are needed to be decoded.

    Return:
        filtered_text : This is the list of all the decoded and processed predictions.

    '''

    # Input length
    input_len = np.ones(shape=pred_label.shape[0]) * pred_label.shape[1]

    # CTC decode
    decode = keras.backend.ctc_decode(pred_label, input_length=input_len, greedy=False, beam_width=5)[0][0][:,:MAX_LABEL_LENGTH]

    # Converting numerics back to their character values
    chars = num_to_char(decode)

    # Join all the characters
    texts = [tf.strings.reduce_join(inputs=char).numpy().decode('UTF-8') for char in chars]

    # Remove the unknown token
    filtered_texts = [text.replace('[UNK]', " ").strip() for text in texts]

    return filtered_texts

def predict_word_images(word_images: List[np.ndarray]) -> List[str]:
    """Predict text from word images"""
    try:
        if not word_images:
            return []
        
        # Preprocess all word images
        processed_images = []
        for word_img in word_images:
            processed_img = preprocess_word_image(word_img)
            if processed_img is not None:
                processed_images.append(processed_img)
        
        if not processed_images:
            return []
        
        # Batch predict
        batch_images = np.array(processed_images)
        predictions = inference_model.predict(batch_images)
        
        # Decode predictions
        decoded_texts = decode_prediction(predictions)
        
        # Apply character corrections
        corrected_texts = [replace_chars(text) for text in decoded_texts]
        
        return corrected_texts
        
    except Exception as e:
        logger.error(f"Error predicting word images: {e}")
        return []

@app.on_event("startup")
async def startup_event():
    """Initialize the model on startup"""
    try:
        load_model_and_setup()
        logger.info("API startup completed successfully")
    except Exception as e:
        logger.error(f"Failed to initialize API: {e}")
        raise

@app.post("/ocr/predict", response_model=Dict[str, Any])
async def predict_text(file: UploadFile = File(...)):
    """
    Extract and recognize text from an uploaded image
    """
    try:
        # Validate file type
        if not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="File must be an image")
        
        # Read uploaded image
        contents = await file.read()
        # Create temporary directory for processing
        with tempfile.TemporaryDirectory() as temp_dir:
            # Save uploaded image
            image_path = os.path.join(temp_dir, f"input_{file.filename}")
            with open(image_path, 'wb') as f:
                f.write(contents)
            
            # Apply CRAFT detection
            craft_output_dir = os.path.join(temp_dir, "craft_output/")
            craft_output_dir = apply_craft_detection(image_path, craft_output_dir)
            # Find bounding box file
            image_basename = os.path.splitext(os.path.basename(image_path))[0]
            bbox_file = os.path.join(craft_output_dir, f"res_{image_basename}.txt")
            if not os.path.exists(bbox_file):
                raise HTTPException(status_code=404, detail="No text detected in image")
            
            # Sort bounding boxes
            sorted_file = sort_bounding_boxes(bbox_file)
            if not sorted_file:
                raise HTTPException(status_code=404, detail="No valid bounding boxes found")
            
            # Extract word images
            word = 0
            output_folder = os.path.join(temp_dir, "extracted_word_images/")
            word = extract_bounding_boxes(image_path, sorted_file, output_folder, word)
            pad_and_resize_images(output_folder)
            test_csv_path = os.path.join(temp_dir, 'testing_data.csv')
            create_csv_from_folder(output_folder, test_csv_path)

            test_csv = pd.read_csv(test_csv_path)
            test_csv['IDENTITY'] = test_csv['IDENTITY'].apply(lambda x: str(x))
            test_csv['FILENAME']  = [output_folder + f"/{filename}" for filename in test_csv['FILENAME']]

            resize_images_in_folder(output_folder)

            df_infer = test_csv

            # Step 1: Sort the dataframe based on values before ';'
            df_infer['before_semicolon'] = df_infer['IDENTITY'].apply(lambda x: int(x.split(';')[0]))
            df_infer['after_semicolon'] = df_infer['IDENTITY'].apply(lambda x: int(x.split(';')[1]))
            sorted_df = df_infer.sort_values(['before_semicolon']).reset_index(drop=True)
            sorted_df.drop(columns=['before_semicolon', 'after_semicolon'], inplace=True)

            sorted_df['IDENTITY'] = sorted_df['IDENTITY'].astype(str)

            sorted_dfs = tf.data.Dataset.from_tensor_slices(
                (np.array(sorted_df['FILENAME'].to_list()), np.array(sorted_df['IDENTITY'].to_list()))
            ).map(encode_single_sample, num_parallel_calls=AUTOTUNE).batch(BATCH_SIZE).prefetch(AUTOTUNE)

            decoded_predictions = decode_pred(inference_model.predict(sorted_dfs))
            pred = sorted_df['IDENTITY'].tolist()
            formatted_output = []

            current_group = None
            i = 0
            for prediction in pred:
                before, after = map(int, prediction.split(';'))

                if current_group is None:
                    current_group = after

                if after != current_group:
                    formatted_output.append('\n')  # Start a new line for the new group
                    current_group = after

                formatted_output.append(decoded_predictions[i] + ' ')
                i += 1

            formatted_output.append('\n')  # Final new line

            full_text = ''.join(formatted_output)


            # if not output_folder:
            #     raise HTTPException(status_code=404, detail="No word images extracted")
            

            # # Predict text from word images
            # predictions = predict_word_images(word_images)
            # # Combine predictions into full text
            # full_text = ' '.join(predictions)
            print(full_text)
            return {
                "status": "success",
                "extracted_text": full_text,
                # "word_count": len(predictions),
                # "words": predictions,
                "message": "Text extracted successfully"
            }
            
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing image: {e}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "model_loaded": inference_model is not None}

@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": "Spanish OCR API",
        "version": "1.0.0",
        "endpoints": {
            "predict": "/ocr/predict",
            "health": "/health"
        }
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)