Upload ocr_server.py

ocr_server.py

@@ -0,0 +1,179 @@
+# ocr_server.py - Clean natural language messages (no Reference number text)
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import base64
+import binascii
+from PIL import Image
+import io
+import torch
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+import re
+
+app = FastAPI()
+
+# Enable CORS for mobile app
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Request model for Base64
+class OCRRequest(BaseModel):
+    image_base64: str
+    filename: str = "image.jpg"
+
+# Load model
+print("Loading OCR model...")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_id = "prithivMLmods/coreOCR-7B-050325-preview"
+
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device).eval()
+
+print(f"Model loaded on {device}")
+
+def extract_numbers(text):
+    numbers = re.findall(r'\d+', text)
+    return ''.join(numbers) if numbers else ""
+
+def base64_to_image(base64_string):
+    if ',' in base64_string and base64_string.startswith('data:'):
+        base64_string = base64_string.split(',', 1)[1]
+    image_bytes = base64.b64decode(base64_string)
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    return image
+
+@app.post("/ocr")
+async def ocr_image(request: OCRRequest):
+    try:
+        # Convert base64 to image
+        image = base64_to_image(request.image_base64)
+        
+        # Run OCR
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": "Extract all numbers from this meter reading. Return only the numbers."},
+            ]
+        }]
+        
+        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(
+            text=[prompt_full],
+            images=[image],
+            return_tensors="pt",
+        ).to(device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, 
+                max_new_tokens=200,
+                temperature=0.1,
+                do_sample=False
+            )
+        
+        result = processor.decode(outputs[0], skip_special_tokens=True)
+        numbers_only = extract_numbers(result)
+        
+        # Clean natural language messages
+        if numbers_only:
+            message = "Meter reading successfully extracted"
+            ref_no = numbers_only[-14:] if len(numbers_only) >= 14 else numbers_only
+        else:
+            message = "No numbers found in the image. Please provide a clear meter reading photo"
+            ref_no = ""
+        
+        return JSONResponse({
+            "success": True,
+            "message": message,
+            "ref_no": ref_no,
+            "numbers": numbers_only
+        })
+        
+    except binascii.Error:
+        return JSONResponse({
+            "success": False,
+            "message": "Invalid image format. Please send a valid Base64 encoded image"
+        }, status_code=400)
+        
+    except Exception as e:
+        error_message = str(e)
+        if "image" in error_message.lower():
+            message = "Could not process the image. Please ensure it's a valid photo of a meter reading"
+        elif "timeout" in error_message.lower():
+            message = "OCR processing timed out. Please try with a smaller or clearer image"
+        else:
+            message = "Failed to process image. Please try again"
+        
+        return JSONResponse({
+            "success": False,
+            "message": message
+        }, status_code=500)
+
+@app.get("/health")
+async def health_check():
+    return {
+        "status": "ok",
+        "model": "coreOCR-7B",
+        "message": "OCR server is running normally"
+    }
+
+@app.post("/ocr-file")
+async def ocr_image_file(file: bytes = None):
+    """Legacy file upload endpoint for testing"""
+    try:
+        image = Image.open(io.BytesIO(file)).convert("RGB")
+        
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": "Extract all numbers from this meter reading. Return only the numbers."},
+            ]
+        }]
+        
+        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(
+            text=[prompt_full],
+            images=[image],
+            return_tensors="pt",
+        ).to(device)
+        
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.1, do_sample=False)
+        
+        result = processor.decode(outputs[0], skip_special_tokens=True)
+        numbers_only = extract_numbers(result)
+        
+        if numbers_only:
+            message = "Meter reading successfully extracted"
+            ref_no = numbers_only[-14:] if len(numbers_only) >= 14 else numbers_only
+        else:
+            message = "No numbers found in the image. Please provide a clear meter reading photo"
+            ref_no = ""
+        
+        return JSONResponse({
+            "success": True,
+            "message": message,
+            "ref_no": ref_no,
+            "numbers": numbers_only
+        })
+    except Exception as e:
+        return JSONResponse({
+            "success": False,
+            "message": "Failed to process image. Please try again"
+        }, status_code=500)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)