Update app.py

app.py

@@ -8,13 +8,15 @@ from transformers import AutoProcessor, AutoModelForImageTextToText, Qwen2_5_VLF
 from reportlab.platypus import SimpleDocTemplate, Paragraph
 from reportlab.lib.styles import getSampleStyleSheet
 from docx import Document
+from gTTS import gTTS
+from jiwer import cer
 
 # ---------------- Models ----------------
 MODEL_PATHS = {
     "Model 1 (Complex handwrittings )": ("prithivMLmods/Qwen2.5-VL-7B-Abliterated-Caption-it", Qwen2_5_VLForConditionalGeneration),
     "Model 2 (simple and scanned handwritting )": ("nanonets/Nanonets-OCR-s", Qwen2_5_VLForConditionalGeneration),
-    "Model 3 (structured handwritting)": ("Emeritus-21/Finetuned-full-HTR-model", AutoModelForImageTextToText),
 }
+# Model 3 has been removed to conserve memory.
 
 MAX_NEW_TOKENS_DEFAULT = 512
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -62,12 +64,10 @@ def _build_inputs(processor, tokenizer, image: Image.Image, prompt: str):
 def _decode_text(model, processor, tokenizer, output_ids, prompt: str):
     try:
         decoded_text = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
-        # Remove the prompt text and any preceding content to ensure only the transcription remains
         prompt_start = decoded_text.find(prompt)
         if prompt_start != -1:
             decoded_text = decoded_text[prompt_start + len(prompt):].strip()
         else:
-            # If prompt is not found, return the stripped decoded text
             decoded_text = decoded_text.strip()
         return decoded_text
     except Exception:
@@ -133,7 +133,6 @@ def save_as_audio(text):
     text = _safe_text(text)
     if not text: return None
     try: 
-        from gTTS import gTTS
         tts = gTTS(text)
         tts.save("output.mp3")
         return "output.mp3"
@@ -141,10 +140,39 @@ def save_as_audio(text):
         print(f"gTTS failed: {e}")
         return None
 
+# ---------------- Metrics Function ----------------
+def calculate_cer_score(ground_truth: str, prediction: str) -> str:
+    """
+    Calculates the Character Error Rate (CER) between two strings.
+    A CER of 0.0 means the prediction is perfect.
+    """
+    if not ground_truth or not prediction:
+        return "Cannot calculate CER: Missing ground truth or prediction."
+    
+    ground_truth_cleaned = " ".join(ground_truth.strip().split())
+    prediction_cleaned = " ".join(prediction.strip().split())
+    
+    error_rate = cer(ground_truth_cleaned, prediction_cleaned)
+    return f"Character Error Rate (CER): {error_rate:.4f}"
+
+# ---------------- Evaluation Orchestration ----------------
+@spaces.GPU
+def perform_evaluation(image: Image.Image, model_name: str, ground_truth: str,
+                       max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float):
+    if image is None or not ground_truth:
+        return "Please upload an image and provide the ground truth.", "N/A"
+    
+    prediction = ocr_image(image, model_name, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty)
+    
+    cer_score = calculate_cer_score(ground_truth, prediction)
+    
+    return prediction, cer_score
+
 # ---------------- Gradio Interface ----------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("## ✍🏾 wilson Handwritten OCR")
     model_choice = gr.Radio(choices=list(MODEL_PATHS.keys()), value=list(MODEL_PATHS.keys())[0], label="Select OCR Model")
+
     with gr.Tab("🖼 Image Inference"):
         query_input = gr.Textbox(label="Custom Prompt (optional)", placeholder="Leave empty for RAW structured output")
         image_input = gr.Image(type="pil", label="Upload / Capture Handwritten Image", sources=["upload", "webcam"])
@@ -168,5 +196,26 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         audio_btn.click(fn=save_as_audio, inputs=[raw_output], outputs=[audio_file])
         clear_btn.click(fn=lambda: ("", None, "", MAX_NEW_TOKENS_DEFAULT, 0.1, 1.0, 0, 1.0), outputs=[raw_output, image_input, query_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty])
 
+    with gr.Tab("📊 Model Evaluation"):
+        gr.Markdown("### 🔍 Evaluate Model Accuracy")
+        eval_image_input = gr.Image(type="pil", label="Upload Image for Evaluation", sources=["upload"])
+        eval_ground_truth = gr.Textbox(label="Ground Truth (Correct Transcription)", lines=10, placeholder="Type or paste the correct text here.")
+        eval_model_output = gr.Textbox(label="Model's Prediction", lines=10, interactive=False, show_copy_button=True)
+        eval_cer_output = gr.Textbox(label="Metrics", interactive=False)
+        
+        with gr.Row():
+            run_evaluation_btn = gr.Button("🚀 Run OCR and Evaluate", variant="primary")
+            clear_evaluation_btn = gr.Button("🧹 Clear")
+            
+        run_evaluation_btn.click(
+            fn=perform_evaluation,
+            inputs=[eval_image_input, model_choice, eval_ground_truth, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+            outputs=[eval_model_output, eval_cer_output]
+        )
+        clear_evaluation_btn.click(
+            fn=lambda: (None, "", "", ""),
+            outputs=[eval_image_input, eval_ground_truth, eval_model_output, eval_cer_output]
+        )
+
 if __name__ == "__main__":
     demo.queue(max_size=50).launch(share=True)