app.py

# -*- coding: utf-8 -*-
# app.py

import os
# Disable Streamlit file watcher to prevent conflicts with PyTorch
os.environ["STREAMLIT_SERVER_ENABLE_FILE_WATCHER"] = "false"

import streamlit as st
import pandas as pd
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import traceback

# Import all necessary configuration values from config.py
from config import (
    IMG_HEIGHT, NUM_CLASSES, BLANK_TOKEN, VOCABULARY, BLANK_TOKEN_SYMBOL,
    TRAIN_CSV_PATH, TEST_CSV_PATH, TRAIN_IMAGES_DIR, TEST_IMAGES_DIR,
    MODEL_SAVE_PATH, BATCH_SIZE, NUM_EPOCHS
)

# Import classes and functions from data_handler_ocr.py and model_ocr.py
from data_handler_ocr import CharIndexer, OCRDataset, ocr_collate_fn, load_ocr_dataframes, create_ocr_dataloaders
from model_ocr import CRNN, train_ocr_model, save_ocr_model, load_ocr_model, ctc_greedy_decode
from utils_ocr import preprocess_user_image_for_ocr, binarize_image, resize_image_for_ocr, normalize_image_for_model


# --- Global Variables ---
ocr_model = None
char_indexer = None 
training_history = None 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- Streamlit App Setup ---
st.set_page_config(layout="wide", page_title="Handwritten Name OCR App",)


st.title("📝 Handwritten Name Recognition (OCR) App")
st.markdown("""
    This application uses a Convolutional Recurrent Neural Network (CRNN) to perform
    Optical Character Recognition (OCR) on handwritten names. You can upload an image
    of a handwritten name for prediction or train a new model using the provided dataset.

    **Note:** Training a robust OCR model can be time-consuming.
""")

# --- Initialize CharIndexer ---
# This initializes char_indexer once when the script starts
char_indexer = CharIndexer(vocabulary_string=VOCABULARY, blank_token_symbol=BLANK_TOKEN_SYMBOL)

# --- Model Loading / Initialization ---
@st.cache_resource # Cache the model to prevent reloading on every rerun
def get_and_load_ocr_model_cached(num_classes, model_path):
    """
    Initializes the OCR model and attempts to load a pre-trained model.
    If no pre-trained model exists, a new model instance is returned.
    """
    model_instance = CRNN(num_classes=num_classes, cnn_output_channels=512, rnn_hidden_size=256, rnn_num_layers=2)
    
    if os.path.exists(model_path):
        st.sidebar.info("Loading pre-trained OCR model...")
        try:
            # Load model to CPU first, then move to device
            model_instance.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
            st.sidebar.success("OCR model loaded successfully!")
        except Exception as e:
            st.sidebar.error(f"Error loading model: {e}. A new model will be initialized.")
            # If loading fails, re-initialize an untrained model
            model_instance = CRNN(num_classes=num_classes, cnn_output_channels=512, rnn_hidden_size=256, rnn_num_layers=2)
    else:
        st.sidebar.warning("No pre-trained OCR model found. Please train a model using the sidebar option.")
    
    return model_instance

# Get the model instance and assign it to the global 'ocr_model'
ocr_model = get_and_load_ocr_model_cached(char_indexer.num_classes, MODEL_SAVE_PATH)
# Ensure the model is on the correct device for inference
ocr_model.to(device)
ocr_model.eval() # Set model to evaluation mode for inference by default


# --- Sidebar for Model Training ---
st.sidebar.header("Train OCR Model")
st.sidebar.write("Click the button below to start training the OCR model.")

# Progress bar and label for training in the sidebar
progress_bar_sidebar = st.sidebar.progress(0)
progress_label_sidebar = st.sidebar.empty()

def update_progress_callback_sidebar(value, text):
    progress_bar_sidebar.progress(int(value * 100))
    progress_label_sidebar.text(text)

if st.sidebar.button("📊 Start Training"):
    progress_bar_sidebar.progress(0)
    progress_label_sidebar.empty()
    st.empty()

    if not os.path.exists(TRAIN_CSV_PATH) or not os.path.isdir(TRAIN_IMAGES_DIR):
        st.sidebar.error(f"Training CSV '{TRAIN_CSV_PATH}' or Images directory '{TRAIN_IMAGES_DIR}' not found!")
    elif not os.path.exists(TEST_CSV_PATH) or not os.path.isdir(TEST_IMAGES_DIR):
        st.sidebar.warning(f"Test CSV '{TEST_CSV_PATH}' or Images directory '{TEST_IMAGES_DIR}' not found. "
                   "Evaluation might be affected or skipped. Please ensure all data paths are correct.")
    else:
        st.sidebar.info(f"Training a new CRNN model for {NUM_EPOCHS} epochs. This will take significant time...")
        
        try:
            train_df, test_df = load_ocr_dataframes(TRAIN_CSV_PATH, TEST_CSV_PATH)
            st.sidebar.success("Training and Test DataFrames loaded successfully.")

            st.sidebar.success(f"CharIndexer initialized with {char_indexer.num_classes} classes.")

            train_loader, test_loader = create_ocr_dataloaders(train_df, test_df, char_indexer, BATCH_SIZE)
            st.sidebar.success("DataLoaders created successfully.")
            
            ocr_model.train() 

            st.sidebar.write("Training in progress... This may take a while.")
            ocr_model, training_history = train_ocr_model(
                model=ocr_model,
                train_loader=train_loader,
                test_loader=test_loader,
                char_indexer=char_indexer,
                epochs=NUM_EPOCHS,
                device=device,
                progress_callback=update_progress_callback_sidebar 
            )
            st.sidebar.success("OCR model training finished!")
            update_progress_callback_sidebar(1.0, "Training complete!")

            os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
            save_ocr_model(ocr_model, MODEL_SAVE_PATH) 
            st.sidebar.success(f"Trained model saved to `{MODEL_SAVE_PATH}`")

        except Exception as e:
            st.sidebar.error(f"An error occurred during training: {e}")
            st.exception(e) 
            update_progress_callback_sidebar(0.0, "Training failed!")

# --- Sidebar for Model Loading ---
st.sidebar.header("Load Pre-trained Model")
st.sidebar.write("If you have a saved model, you can load it here instead of training.")

if st.sidebar.button("💾 Load Model"):
    if os.path.exists(MODEL_SAVE_PATH):
        try:
            loaded_model = CRNN(num_classes=char_indexer.num_classes)
            load_ocr_model(loaded_model, MODEL_SAVE_PATH)
            loaded_model.to(device)
            
            st.sidebar.success(f"Model loaded successfully from `{MODEL_SAVE_PATH}`")
        except Exception as e:
            st.sidebar.error(f"Error loading model: {e}")
            st.exception(e) 
    else:
        st.sidebar.warning(f"No model found at `{MODEL_SAVE_PATH}`. Please train a model first or check the path.")

# --- Main Content: Prediction Section and Training History  ---

# Display training history chart
if training_history:
    st.subheader("Training History Plots") 
    history_df = pd.DataFrame({
        'Epoch': range(1, len(training_history['train_loss']) + 1),
        'Train Loss': training_history['train_loss'],
        'Test Loss': training_history['test_loss'],
        'Test CER (%)': [cer * 100 for cer in training_history['test_cer']],
        'Test Exact Match Accuracy (%)': [acc * 100 for acc in training_history['test_exact_match_accuracy']]
    })

    st.markdown("**Loss over Epochs**")
    st.line_chart(history_df.set_index('Epoch')[['Train Loss', 'Test Loss']])
    st.caption("Lower loss indicates better model performance.")

    st.markdown("**Character Error Rate (CER) over Epochs**")
    st.line_chart(history_df.set_index('Epoch')[['Test CER (%)']])
    st.caption("Lower CER indicates fewer character errors (0% is perfect).")

    st.markdown("**Exact Match Accuracy over Epochs**")
    st.line_chart(history_df.set_index('Epoch')[['Test Exact Match Accuracy (%)']])
    st.caption("Higher exact match accuracy indicates more perfectly recognized names.")

    st.markdown("**Performance Metrics over Epochs (CER vs. Exact Match Accuracy)**")
    st.line_chart(history_df.set_index('Epoch')[['Test CER (%)', 'Test Exact Match Accuracy (%)']])
    st.caption("CER should decrease, Accuracy should increase.")
    st.write("---") # Separator after charts


# Predict on a New Image

if ocr_model is None:
    st.warning("Please train or load a model before attempting prediction.")
else:
    uploaded_file = st.file_uploader("🖼️ Choose an image...", type=["png", "jpg", "jpeg", "jfif"])

    if uploaded_file is not None:
        try:
            image_pil = Image.open(uploaded_file).convert('L')
            st.image(image_pil, caption="Uploaded Image", use_container_width=True) 
            st.write("---")
            st.write("Processing and Recognizing...")

            processed_image_tensor = preprocess_user_image_for_ocr(image_pil, IMG_HEIGHT).to(device)
            
            ocr_model.eval()
            with torch.no_grad():
                output = ocr_model(processed_image_tensor)
            
            predicted_texts = ctc_greedy_decode(output, char_indexer)
            predicted_text = predicted_texts[0]

            st.success(f"Recognized Text: **{predicted_text}**")

        except Exception as e:
            st.error(f"Error processing image or recognizing text: {e}")
            st.info("💡 **Tips for best results:**\n"
                    "- Ensure the handwritten text is clear and on a clean background.\n"
                    "- Only include one name/word per image.\n"
                    "- The model is trained on specific characters. Unusual symbols might not be recognized.")
            st.exception(e)

st.markdown("""
    ---
    *Built using Streamlit, PyTorch, OpenCV, and EditDistance ©2025 by MFT*
    """)