Update app.py

app.py

@@ -4,7 +4,7 @@ from flask import Flask, request, jsonify
 from flask_cors import CORS
 from werkzeug.utils import secure_filename
 import torch
-from transformers import DonutProcessor, VisionEncoderDecoderModel
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
 import cv2
 import numpy as np
@@ -21,13 +21,14 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max upload
 # Create uploads directory if it doesn't exist
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
-# Load OCR model - Microsoft's Donut model
+# Load OCR model - Microsoft's TrOCR model
 print("Loading OCR model...")
-processor = DonutProcessor.from_pretrained("microsoft/donut-base", cache_dir="/huggingface_cache")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/donut-base", cache_dir="/huggingface_cache")
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten", cache_dir="/huggingface_cache")
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten", cache_dir="/huggingface_cache")
 
 # Move model to GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 model.to(device)
 
 def allowed_file(filename):
@@ -60,27 +61,16 @@ def perform_ocr(image_path):
     pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
     
     # Generate text
-    task_prompt = "<s_ocr>"
-    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
-    
-    outputs = model.generate(
+    generated_ids = model.generate(
         pixel_values,
-        decoder_input_ids=decoder_input_ids,
-        max_length=model.decoder.config.max_position_embeddings,
-        early_stopping=True,
-        pad_token_id=processor.tokenizer.pad_token_id,
-        eos_token_id=processor.tokenizer.eos_token_id,
-        use_cache=True,
+        max_length=64,
         num_beams=5,
-        bad_words_ids=[[processor.tokenizer.unk_token_id]],
-        return_dict_in_generate=True,
+        early_stopping=True
     )
     
     # Decode generated text
-    sequence = processor.batch_decode(outputs.sequences)[0]
-    sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-    sequence = sequence.replace("<s>", "").replace("</s>", "").replace("<s_ocr>", "").replace("</s_ocr>", "")
-    return sequence.strip()
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text.strip()
 
 @app.route('/ocr', methods=['POST'])
 def ocr():
@@ -107,14 +97,16 @@ def ocr():
             # Perform OCR
             extracted_text = perform_ocr(file_path)
             
-            # Clean up the file if needed
-            # os.remove(file_path)
+            # Clean up the file after processing
+            os.remove(file_path)
             
             return jsonify({
                 'success': True,
                 'text': extracted_text
             })
         except Exception as e:
+            # Log the error
+            print(f"Error processing image: {str(e)}")
             return jsonify({
                 'success': False,
                 'error': str(e)