Add internal ollama parsing to citizenship ocr

Dockerfile

@@ -57,6 +57,7 @@ ENV DOCTR_CACHE_DIR=/app/.cache/doctr
 # Copy the application code
 COPY main.py .
 COPY utils.py ./
+COPY data_models.py ./
 COPY models /app/models
 
 # Expose the port FastAPI will run on

data_models.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+
+class Citizenship(BaseModel):
+    citizenship_no: str
+    name: str
+    gender: str
+    permanent_district: str
+    permanent_nagarpalika: str
+    temporary_district: str
+    temporary_nagarpalika: str
+    date_of_birth: str
+    father_name: str
+    father_address: str
+    mother_name: str
+    mother_address: str
+    wife_name: str
+    wife_address: str

main.py

@@ -71,7 +71,7 @@ from pydantic import BaseModel
 import shutil
 
 # Import from optimized utils
-from utils import dev_number, roman_number, dev_letter, roman_letter, predict_ne, ocr_citizenship_utils
+from utils import dev_number, roman_number, dev_letter, roman_letter, predict_ne, perform_citizenship_ocr
 
 app = FastAPI(
     title="OCR API",
@@ -198,10 +198,9 @@ async def classify_ne(image: UploadFile = File(...)):
 async def ocr_citizenship(image: UploadFile = File(...)):
     """OCR the provided Nepali Citizenship card"""
     image_path  = await save_upload_file_tmp(image)
-    prediction = ocr_citizenship_utils(
+    prediction = perform_citizenship_ocr(
         image_path=image_path,
-    )
-    
+    )    
     return JSONResponse(content=prediction)
 # Health check endpoint
 @app.get("/health")

requirements.txt

@@ -1,4 +1,4 @@
-python-doctr[torch,viz]
+python-doctr[torch,viz]==0.11.0
 torch
 torchvision
 numpy
@@ -9,5 +9,5 @@ uvicorn
 pydantic
 python-multipart
 scikit-learn==1.6.1
-opencv-python
+opencv-python==4.11.0.86
 surya-ocr==0.13.1

utils.py

@@ -6,12 +6,16 @@ from PIL import Image
 # from functools import lru_cache
 from torchvision import models
 from typing import List
+from fastapi import HTTPException
+from data_models import Citizenship
+import json
 import torchvision.transforms as transforms
 import torch
 import torch.nn as nn
 import numpy as np
 import cv2
 import regex as re
+import requests
 # import os
 import pickle
 
@@ -323,7 +327,7 @@ def merge_boxes_same_line(boxes, y_thresh=5, x_thresh=60):
 
     return np.array(merged)
 
-def ocr_citizenship_utils(image_path: str) -> List[List[str]]:
+def ocr_citizenship(image_path: str) -> List[List[str]]:
     doctr_detector, surya_recognition_predictor, surya_detection_predictor = initialize_detector()
     page = cv2.imread(image_path)
     page = cv2.convertScaleAbs(page, alpha=1.5, beta=0)
@@ -368,4 +372,45 @@ def ocr_citizenship_utils(image_path: str) -> List[List[str]]:
         line_result.append(text_combo)
         prev_y = boxes[1] 
     
-    return full_result
+    return full_result
+
+PARSE_PROMPT = "You are a parsing agent. Your task is to generate a json response from the given text corpus."
+def create_local_model(message, base_model):
+    try:
+        ollama_endpoint = "api/chat"
+        url = f"https://aioverlords-amnil-internal-ollama.hf.space/proxy/{ollama_endpoint}"
+
+        # Data to send in the POST request
+        data = {
+            "data": {
+                "model": "aisingapore/Llama-SEA-LION-v3-8B-IT",
+                "messages": message,
+                "stream": False,
+                "format": base_model.model_json_schema()
+            }
+        }
+
+        response = requests.post(url, json=data)
+        # Check the response
+        if response.status_code == 200:
+            print(f"Request Success:", response.json())
+            return json.loads(response.json()["message"]["content"])
+            # return response.json()
+        else:
+            print(f"Request Error:", response.status_code, response.text)
+            raise HTTPException(status_code=response.status_code, detail=response.text)
+    except HTTPException as http_exec:
+        raise http_exec
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    
+def perform_citizenship_ocr(image_path):
+    try:        
+        unparsed_result = ocr_citizenship(image_path)
+        message = [
+            {"role": "system", "content": PARSE_PROMPT},
+            {"role": "user", "content": f"Given Text: \n{unparsed_result}"},
+        ]
+        return create_local_model(message, Citizenship)        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))