Create app.py

app.py

@@ -0,0 +1,400 @@
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import base64
+import io
+import os
+from PIL import Image
+import logging
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import torch
+import easyocr
+import numpy as np
+import threading
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = Flask(__name__)
+CORS(app)
+
+# Global variables for models
+trocr_processor = None
+trocr_model = None
+easyocr_reader = None
+models_loaded = False
+loading_lock = threading.Lock()
+
+def initialize_models():
+    """Initialize OCR models"""
+    global trocr_processor, trocr_model, easyocr_reader, models_loaded
+    
+    if models_loaded:
+        return
+    
+    with loading_lock:
+        if models_loaded:  # Double-check after acquiring lock
+            return
+            
+        try:
+            logger.info("Starting model initialization...")
+            
+            # Set cache directory
+            cache_dir = os.environ.get('TRANSFORMERS_CACHE', '/app/.cache/huggingface')
+            os.makedirs(cache_dir, exist_ok=True)
+            
+            # Initialize TrOCR for handwritten text (Microsoft's model)
+            logger.info("Loading TrOCR model for handwritten text...")
+            trocr_processor = TrOCRProcessor.from_pretrained(
+                "microsoft/trocr-base-handwritten",
+                cache_dir=cache_dir
+            )
+            trocr_model = VisionEncoderDecoderModel.from_pretrained(
+                "microsoft/trocr-base-handwritten",
+                cache_dir=cache_dir
+            )
+            
+            # Initialize EasyOCR for printed text
+            logger.info("Loading EasyOCR for printed text...")
+            easyocr_reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
+            
+            models_loaded = True
+            logger.info("All models loaded successfully!")
+            
+        except Exception as e:
+            logger.error(f"Error loading models: {str(e)}")
+            models_loaded = False
+            raise e
+
+def ensure_models_loaded():
+    """Ensure models are loaded before processing"""
+    if not models_loaded:
+        initialize_models()
+
+def preprocess_image(image):
+    """Preprocess image for better OCR results"""
+    # Convert to RGB if needed
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
+    # Resize if image is too large
+    max_size = 1024
+    if max(image.size) > max_size:
+        ratio = max_size / max(image.size)
+        new_size = tuple(int(dim * ratio) for dim in image.size)
+        image = image.resize(new_size, Image.Resampling.LANCZOS)
+    
+    return image
+
+def extract_text_trocr(image):
+    """Extract text using TrOCR (good for handwritten text)"""
+    try:
+        ensure_models_loaded()
+        if not trocr_processor or not trocr_model:
+            return ""
+            
+        # Preprocess image
+        image = preprocess_image(image)
+        
+        # Generate pixel values
+        pixel_values = trocr_processor(image, return_tensors="pt").pixel_values
+        
+        # Generate text
+        generated_ids = trocr_model.generate(pixel_values)
+        generated_text = trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
+        return generated_text.strip()
+    except Exception as e:
+        logger.error(f"TrOCR error: {str(e)}")
+        return ""
+
+def extract_text_easyocr(image):
+    """Extract text using EasyOCR (good for printed text)"""
+    try:
+        ensure_models_loaded()
+        if not easyocr_reader:
+            return ""
+            
+        # Convert PIL image to numpy array
+        image_np = np.array(preprocess_image(image))
+        
+        # Extract text
+        results = easyocr_reader.readtext(image_np, detail=0)
+        
+        # Join all detected text
+        extracted_text = ' '.join(results)
+        return extracted_text.strip()
+    except Exception as e:
+        logger.error(f"EasyOCR error: {str(e)}")
+        return ""
+
+def process_image_ocr(image, ocr_type="auto"):
+    """Process image with specified OCR method"""
+    results = {}
+    
+    if ocr_type in ["auto", "handwritten", "trocr"]:
+        trocr_text = extract_text_trocr(image)
+        results["trocr"] = trocr_text
+    
+    if ocr_type in ["auto", "printed", "easyocr"]:
+        easyocr_text = extract_text_easyocr(image)
+        results["easyocr"] = easyocr_text
+    
+    # For auto mode, return the longer result or combine both
+    if ocr_type == "auto":
+        trocr_len = len(results.get("trocr", ""))
+        easyocr_len = len(results.get("easyocr", ""))
+        
+        if trocr_len > 0 and easyocr_len > 0:
+            # If both have results, combine them intelligently
+            if abs(trocr_len - easyocr_len) / max(trocr_len, easyocr_len) < 0.3:
+                # If lengths are similar, prefer EasyOCR for printed text
+                results["final"] = results["easyocr"]
+            else:
+                # Use the longer result
+                results["final"] = results["trocr"] if trocr_len > easyocr_len else results["easyocr"]
+        elif trocr_len > 0:
+            results["final"] = results["trocr"]
+        elif easyocr_len > 0:
+            results["final"] = results["easyocr"]
+        else:
+            results["final"] = ""
+    else:
+        # Return the specific model result
+        results["final"] = results.get(ocr_type.replace("handwritten", "trocr").replace("printed", "easyocr"), "")
+    
+    return results
+
+@app.route('/')
+def home():
+    """Root endpoint"""
+    return jsonify({
+        "service": "OCR Backend",
+        "status": "running",
+        "version": "1.0.0",
+        "models_loaded": models_loaded,
+        "endpoints": {
+            "health": "/health",
+            "ocr": "/ocr (POST)",
+            "batch_ocr": "/ocr/batch (POST)",
+            "models_info": "/models/info (GET)"
+        },
+        "supported_formats": ["PNG", "JPEG", "JPG", "BMP", "TIFF"],
+        "ocr_types": ["auto", "handwritten", "printed"]
+    })
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint"""
+    return jsonify({
+        "status": "healthy", 
+        "models_loaded": models_loaded,
+        "service": "OCR Backend"
+    })
+
+@app.route('/ocr', methods=['POST'])
+def ocr_endpoint():
+    """Main OCR endpoint"""
+    try:
+        # Ensure models are loaded
+        ensure_models_loaded()
+        
+        # Check if image is provided
+        if 'image' not in request.files and not request.is_json:
+            return jsonify({"error": "No image provided. Use 'image' field for file upload or JSON with 'image_base64'"}), 400
+        
+        if request.is_json and 'image_base64' not in request.json:
+            return jsonify({"error": "No 'image_base64' field found in JSON"}), 400
+        
+        # Get OCR type preference
+        if request.is_json:
+            ocr_type = request.json.get('type', 'auto')
+        else:
+            ocr_type = request.form.get('type', 'auto')
+        
+        # Validate ocr_type
+        if ocr_type not in ['auto', 'handwritten', 'printed', 'trocr', 'easyocr']:
+            return jsonify({"error": "Invalid OCR type. Use: auto, handwritten, printed"}), 400
+        
+        # Load image
+        if 'image' in request.files:
+            # File upload
+            image_file = request.files['image']
+            if image_file.filename == '':
+                return jsonify({"error": "No file selected"}), 400
+            image = Image.open(image_file.stream)
+        else:
+            # Base64 image
+            image_data = request.json['image_base64']
+            if image_data.startswith('data:image'):
+                # Remove data URL prefix
+                image_data = image_data.split(',')[1]
+            
+            try:
+                # Decode base64
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(io.BytesIO(image_bytes))
+            except Exception as e:
+                return jsonify({"error": f"Invalid base64 image data: {str(e)}"}), 400
+        
+        # Process image
+        results = process_image_ocr(image, ocr_type)
+        
+        response = {
+            "success": True,
+            "text": results["final"],
+            "type_used": ocr_type,
+            "character_count": len(results["final"]),
+            "details": {
+                "trocr_result": results.get("trocr", ""),
+                "easyocr_result": results.get("easyocr", "")
+            } if ocr_type == "auto" else {}
+        }
+        
+        return jsonify(response)
+        
+    except Exception as e:
+        logger.error(f"OCR processing error: {str(e)}")
+        return jsonify({"error": str(e), "success": False}), 500
+
+@app.route('/ocr/batch', methods=['POST'])
+def batch_ocr_endpoint():
+    """Batch OCR endpoint for multiple images"""
+    try:
+        # Ensure models are loaded
+        ensure_models_loaded()
+        
+        if 'images' not in request.files:
+            return jsonify({"error": "No images provided. Use 'images' field for multiple file upload"}), 400
+        
+        images = request.files.getlist('images')
+        if not images or len(images) == 0:
+            return jsonify({"error": "No images found in request"}), 400
+            
+        ocr_type = request.form.get('type', 'auto')
+        
+        # Validate ocr_type
+        if ocr_type not in ['auto', 'handwritten', 'printed', 'trocr', 'easyocr']:
+            return jsonify({"error": "Invalid OCR type. Use: auto, handwritten, printed"}), 400
+        
+        results = []
+        for i, image_file in enumerate(images):
+            try:
+                if image_file.filename == '':
+                    results.append({
+                        "index": i,
+                        "filename": "empty_file",
+                        "error": "Empty filename",
+                        "success": False
+                    })
+                    continue
+                    
+                image = Image.open(image_file.stream)
+                ocr_results = process_image_ocr(image, ocr_type)
+                
+                results.append({
+                    "index": i,
+                    "filename": image_file.filename,
+                    "text": ocr_results["final"],
+                    "character_count": len(ocr_results["final"]),
+                    "success": True
+                })
+            except Exception as e:
+                results.append({
+                    "index": i,
+                    "filename": image_file.filename if hasattr(image_file, 'filename') else f"image_{i}",
+                    "error": str(e),
+                    "success": False
+                })
+        
+        successful_count = sum(1 for r in results if r["success"])
+        
+        return jsonify({
+            "success": True,
+            "results": results,
+            "total_processed": len(results),
+            "successful": successful_count,
+            "failed": len(results) - successful_count,
+            "type_used": ocr_type
+        })
+        
+    except Exception as e:
+        logger.error(f"Batch OCR error: {str(e)}")
+        return jsonify({"error": str(e), "success": False}), 500
+
+@app.route('/models/info', methods=['GET'])
+def models_info():
+    """Get information about loaded models"""
+    return jsonify({
+        "models": {
+            "trocr": {
+                "name": "microsoft/trocr-base-handwritten",
+                "description": "Handwritten text recognition using Transformer-based OCR",
+                "loaded": trocr_model is not None and trocr_processor is not None,
+                "best_for": "Handwritten text, notes, forms"
+            },
+            "easyocr": {
+                "name": "EasyOCR",
+                "description": "Printed text recognition with CRAFT + CRNN",
+                "loaded": easyocr_reader is not None,
+                "best_for": "Printed text, documents, signs, books"
+            }
+        },
+        "supported_types": ["auto", "handwritten", "printed"],
+        "supported_formats": ["PNG", "JPEG", "JPG", "BMP", "TIFF"],
+        "cache_directory": os.environ.get('TRANSFORMERS_CACHE', '/app/.cache/huggingface'),
+        "gpu_available": torch.cuda.is_available(),
+        "models_loaded": models_loaded
+    })
+
+@app.route('/models/load', methods=['POST'])
+def load_models():
+    """Manually trigger model loading"""
+    try:
+        if models_loaded:
+            return jsonify({"message": "Models already loaded", "success": True})
+        
+        initialize_models()
+        return jsonify({"message": "Models loaded successfully", "success": True})
+    except Exception as e:
+        return jsonify({"error": str(e), "success": False}), 500
+
+@app.errorhandler(404)
+def not_found(error):
+    return jsonify({
+        "error": "Endpoint not found",
+        "available_endpoints": {
+            "GET /": "Service information",
+            "GET /health": "Health check",
+            "POST /ocr": "Single image OCR",
+            "POST /ocr/batch": "Batch image OCR",
+            "GET /models/info": "Model information",
+            "POST /models/load": "Load models manually"
+        }
+    }), 404
+
+@app.errorhandler(500)
+def internal_error(error):
+    return jsonify({
+        "error": "Internal server error",
+        "message": "Please check the server logs for more details"
+    }), 500
+
+# Initialize models when running with gunicorn
+if __name__ != '__main__':
+    logger.info("Starting OCR service with gunicorn...")
+    # Don't initialize models here - let them load lazily on first request
+    # This prevents startup failures due to model loading issues
+
+if __name__ == '__main__':
+    logger.info("Starting OCR service in development mode...")
+    try:
+        # Try to initialize models, but don't fail if it doesn't work
+        initialize_models()
+    except Exception as e:
+        logger.warning(f"Could not initialize models on startup: {e}")
+        logger.info("Models will be loaded on first request")
+    
+    # Run the app
+    port = int(os.environ.get('PORT', 5000))
+    app.run(host='0.0.0.0', port=port, debug=False)