OCR API

Dockerfile

@@ -1,9 +1,13 @@
-FROM python:3.7-slim
+FROM python:3.7
 
 WORKDIR /code
 
 COPY requirements-fastapi.txt ./
 
+# Installing libGL
+RUN apt-get update && apt-get install -y \
+    libgl1-mesa-dev
+
 RUN pip install --no-cache-dir --upgrade -r /code/requirements-fastapi.txt
 
 RUN useradd -m -u 1000 user

config.py

@@ -4,7 +4,9 @@ import os
 
 class Settings(BaseSettings):
     huggingface_key: str = os.environ.get("huggingface_key")
+    sparrow_key: str = os.environ.get("sparrow_key")
     dataset_name: str = "katanaml-org/invoices-donut-data-v1"
+    ocr_stats_file: str = "data/ocr_stats.json"
 
 
 settings = Settings()

data/ocr_stats.json

@@ -0,0 +1 @@
+[[0.0, "receipt_00001.png", "2023-05-23 10:55:43"], [19.22510600090027, "receipt_00001.png", "2023-05-23 11:11:02"], [4.0531158447265625e-06, "receipt_00001.png", "2023-05-23 11:11:10"], [3.0994415283203125e-06, "receipt_00001.png", "2023-05-23 11:11:11"], [3.0994415283203125e-06, "receipt_00001.png", "2023-05-23 11:11:12"], [2.86102294921875e-06, "receipt_00001.png", "2023-05-23 11:11:13"], [3.0994415283203125e-06, "receipt_00001.png", "2023-05-23 11:11:13"], [3.653481960296631, "receipt_00001.png", "2023-05-23 11:32:48"], [8.929341077804565, "receipt_00001.png", "2023-05-23 11:34:52"], [3.5088820457458496, "receipt_00001.png", "2023-05-23 16:32:17"], [2.863774061203003, "receipt_00001.png", "2023-05-23 16:32:56"], [4.174198150634766, "inout-20211211_001.jpg", "2023-05-23 16:38:33"], [4.616858243942261, "inout-20211211_001.jpg", "2023-05-23 16:39:28"], [4.6479880809783936, "inout-20211211_001.jpg", "2023-05-23 16:47:27"], [4.756654262542725, "inout-20211211_001.jpg", "2023-05-23 22:07:14"], [5.704661130905151, "wholefoods-20211211_005.jpg", "2023-05-23 22:16:38"], [6.363792896270752, "wholefoods-20211211_005.jpg", "2023-05-23 22:24:13"], [6.582294940948486, "cvs-20211211_009.jpg", "2023-05-23 22:43:41"], [8.032721757888794, "oldnavy-20211211_015.jpg", "2023-05-23 22:45:58"], [6.35598611831665, "ross-20211211_010.jpg", "2023-05-23 22:47:50"], [7.241703987121582, "ross-20211211_010.jpg", "2023-05-24 11:03:57"], [6.259234189987183, "ross-20211211_010.pdf", "2023-05-24 11:04:26"], [7.275213956832886, "ross-20211211_010.pdf", "2023-05-24 11:05:25"], [5.848371982574463, "invoice_10.jpg", "2023-05-24 11:06:21"], [6.028747081756592, "invoice_10.jpg", "2023-05-24 11:12:14"], [6.5253260135650635, "cvs-20211211_009.jpg", "2023-05-24 11:29:53"], [6.507750988006592, "cvs-20211211_009.jpg", "2023-05-24 12:55:14"], [2.864002227783203, "receipt_00001.png", "2023-05-24 12:55:30"], [2.9030818939208984, "receipt_00001.png", "2023-05-24 12:55:40"], [5.672614097595215, "wholefoods-20211211_005.jpg", "2023-05-24 12:56:13"], [5.712976932525635, "wholefoods-20211211_005.pdf", "2023-05-24 12:56:29"], [5.984729051589966, "invoice_10.jpg", "2023-05-24 13:00:23"], [7.3337507247924805, "bestbuy-20211211_006.pdf", "2023-05-24 13:01:13"], [4.676954984664917, "inout-20211211_001.jpg", "2023-05-24 21:09:53"], [3.9793169498443604, "inout-20211211_001.jpg", "2023-05-24 22:01:12"], [4.716302871704102, "inout-20211211_001.jpg", "2023-05-24 22:07:19"], [4.611649990081787, "inout-20211211_001.jpg", "2023-05-24 22:11:00"], [5.18176007270813, "inout-20211211_001.jpg", "2023-05-24 22:12:26"], [4.76771092414856, "inout-20211211_001.jpg", "2023-05-25 10:00:11"], [4.62838888168335, "inout-20211211_001.jpg", "2023-05-25 10:12:36"], [4.6390650272369385, "inout-20211211_001.jpg", "2023-05-25 10:35:31"], [4.605455160140991, "inout-20211211_001.jpg", "2023-05-25 10:36:59"], [4.541555881500244, "inout-20211211_001.jpg", "2023-05-25 10:37:41"], [4.652244806289673, "inout-20211211_001.jpg", "2023-05-25 10:38:09"], [3.947613000869751, "inout-20211211_001.jpg", "2023-05-25 10:58:54"], [4.597126245498657, "inout-20211211_001.jpg", "2023-05-25 11:00:03"], [4.6871421337127686, "inout-20211211_001.jpg", "2023-05-25 11:02:44"], [4.579195976257324, "inout-20211211_001.jpg", "2023-05-25 11:32:11"], [4.734511137008667, "inout-20211211_001.jpg", "2023-05-25 11:33:03"], [4.602473258972168, "inout-20211211_001.jpg", "2023-05-25 11:44:11"], [4.563000202178955, "inout-20211211_001.jpg", "2023-05-25 11:47:35"], [4.576035022735596, "inout-20211211_001.jpg", "2023-05-25 11:49:55"], [4.860241889953613, "inout-20211211_001.jpg", "2023-05-25 11:53:19"], [4.693282127380371, "inout-20211211_001.jpg", "2023-05-25 11:56:00"], [4.5564610958099365, "inout-20211211_001.jpg", "2023-05-25 16:02:52"], [5.022596836090088, "inout-20211211_001.jpg", "2023-05-25 16:03:47"], [4.650119781494141, "inout-20211211_001.jpg", "2023-05-25 16:27:35"], [6.16159200668335, "inout-20211211_001.jpg", "2023-05-30 22:15:29"], [9.421452045440674, "../docs/models/donut/data/img/test/invoice_2.jpg", "2023-06-07 21:03:34"]]

data/result.json

@@ -0,0 +1,53 @@
+['YOUR GUEST NUMBER IS']
+['43']
+['IN-N-OUT BURGER LINQ']
+['320 6 2166 6301']
+['Cashier: SERJI0 SA']
+['Check : 43']
+['TRANS #: 6301']
+['1 Db1-Dbl']
+['5.25']
+['+ Onion']
+['1 Fry']
+['2.35']
+['1 Med Soft Drink']
+['2.15']
+['COUNTER-Eat In']
+['9.75']
+['TAX 8.375%']
+['.82']
+['Amount Due']
+['$10.57']
+['Tender MasterCard']
+['$10.57']
+['Change']
+['$.00']
+['CHARGE DETAIL']
+['SALE']
+['Card Type:']
+['Mastercard']
+['Account :']
+['************5562 R']
+['Auth Code:']
+['NDTQU8']
+['Trans #:']
+['6301']
+['Auth Ref :']
+['2015517078']
+['AUTH AMT :']
+['$10.57']
+['AID:']
+['A0000000041010']
+['TVR :']
+['0000008001']
+['TSI:']
+['0000']
+['App Name:']
+['Debit MasterCard']
+['ARQC:']
+['ADCF5208793B7BD6']
+['THANK YOU!']
+['Quest ions/Comments: Cal1 800-786-1 :10']
+['L1 T6']
+['9:21 PM']
+['2021-11-30']

endpoints.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from routers import dataset
+from routers import ocr
 
 app = FastAPI(openapi_url="/api/v1/sparrow-data/openapi.json", docs_url="/api/v1/sparrow-data/docs")
 
@@ -13,6 +14,7 @@ app.add_middleware(
 )
 
 app.include_router(dataset.router, prefix="/api-dataset/v1/sparrow-data", tags=["Dataset"])
+app.include_router(ocr.router, prefix="/api-ocr/v1/sparrow-data", tags=["OCR"])
 
 
 @app.get("/")

requirements-fastapi.txt

@@ -1,4 +1,8 @@
+pdf2image==1.16.2
 datasets==2.10.1
 Pillow==9.5.0
-fastapi==0.95.2
+paddlepaddle==2.4.2
+paddleocr==2.6.1.3
+fastapi==0.96.0
+python-multipart
 uvicorn[standard]

routers/dataset.py

@@ -7,11 +7,11 @@ from io import BytesIO
 from PIL import Image
 import base64
 from config import settings
+from huggingface_hub import login
 
 
 router = APIRouter()
 
-from huggingface_hub import login
 login(settings.huggingface_key)
 
 class ImageResponse(BaseModel):
@@ -24,6 +24,7 @@ def encode_pil_image(image: Image) -> str:
     img_data = buffer.getvalue()
     return base64.b64encode(img_data).decode('utf-8')
 
+
 @router.get("/dataset_info")
 async def get_dataset_info():
     dataset = load_dataset(settings.dataset_name)

routers/ocr.py

@@ -0,0 +1,122 @@
+from fastapi import APIRouter, File, UploadFile, Form
+from typing import Optional
+from config import settings
+from PIL import Image
+import urllib.request
+from io import BytesIO
+import utils
+import os
+import time
+from functools import lru_cache
+from paddleocr import PaddleOCR
+from pdf2image import convert_from_bytes
+import io
+import json
+from routers.ocr_utils import merge_data
+
+
+router = APIRouter()
+
+@lru_cache(maxsize=1)
+def load_ocr_model():
+    model = PaddleOCR(use_angle_cls=True, lang='en')
+    return model
+
+
+def invoke_ocr(doc, content_type):
+    worker_pid = os.getpid()
+    print(f"Handling OCR request with worker PID: {worker_pid}")
+    start_time = time.time()
+
+    model = load_ocr_model()
+
+    bytes_img = io.BytesIO()
+
+    format_img = "JPEG"
+    if content_type == "image/png":
+        format_img = "PNG"
+
+    doc.save(bytes_img, format=format_img)
+    bytes_data = bytes_img.getvalue()
+    bytes_img.close()
+
+    result = model.ocr(bytes_data, cls=True)
+
+    values = []
+    for idx in range(len(result)):
+        res = result[idx]
+        for line in res:
+            values.append(line)
+
+    values = merge_data(values)
+
+    end_time = time.time()
+    processing_time = end_time - start_time
+    print(f"OCR done, worker PID: {worker_pid}")
+
+    return values, processing_time
+
+@router.post("/ocr")
+async def run_ocr(file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None),
+                  sparrow_key: str = Form(None)):
+
+    if sparrow_key != settings.sparrow_key:
+        return {"error": "Invalid Sparrow key."}
+
+    result = []
+    if file:
+        if file.content_type in ["image/jpeg", "image/jpg", "image/png"]:
+            doc = Image.open(BytesIO(await file.read()))
+        elif file.content_type == "application/pdf":
+            pdf_bytes = await file.read()
+            pages = convert_from_bytes(pdf_bytes, 300)
+            doc = pages[0]
+        else:
+            return {"error": "Invalid file type. Only JPG/PNG images and PDF are allowed."}
+
+        result, processing_time = invoke_ocr(doc, file.content_type)
+
+        utils.log_stats(settings.ocr_stats_file, [processing_time, file.filename])
+        print(f"Processing time OCR: {processing_time:.2f} seconds")
+    elif image_url:
+        # test image url: https://raw.githubusercontent.com/katanaml/sparrow/main/sparrow-data/docs/input/invoices/processed/images/invoice_10.jpg
+        # test PDF: https://raw.githubusercontent.com/katanaml/sparrow/main/sparrow-data/docs/input/receipts/2021/us/bestbuy-20211211_006.pdf
+        with urllib.request.urlopen(image_url) as response:
+            content_type = response.info().get_content_type()
+
+            if content_type in ["image/jpeg", "image/jpg", "image/png"]:
+                doc = Image.open(BytesIO(response.read()))
+            elif content_type == "application/octet-stream":
+                pdf_bytes = response.read()
+                pages = convert_from_bytes(pdf_bytes, 300)
+                doc = pages[0]
+            else:
+                return {"error": "Invalid file type. Only JPG/PNG images and PDF are allowed."}
+
+        result, processing_time = invoke_ocr(doc, content_type)
+
+        # parse file name from url
+        file_name = image_url.split("/")[-1]
+        utils.log_stats(settings.ocr_stats_file, [processing_time, file_name])
+        print(f"Processing time OCR: {processing_time:.2f} seconds")
+    else:
+        result = {"info": "No input provided"}
+
+    return result
+
+
+@router.get("/statistics")
+async def get_statistics():
+    file_path = settings.ocr_stats_file
+
+    # Check if the file exists, and read its content
+    if os.path.exists(file_path):
+        with open(file_path, 'r') as file:
+            try:
+                content = json.load(file)
+            except json.JSONDecodeError:
+                content = []
+    else:
+        content = []
+
+    return content

routers/ocr_utils.py

@@ -0,0 +1,7 @@
+def merge_data(values):
+    data = []
+    for idx in range(len(values)):
+        data.append([values[idx][1][0]])
+        # print(data[idx])
+
+    return data

utils.py

@@ -0,0 +1,29 @@
+import json
+import os
+from datetime import datetime
+
+
+def log_stats(file_path, new_data):
+    # Check if the file exists, and read its content
+    if os.path.exists(file_path):
+        with open(file_path, 'r') as file:
+            try:
+                content = json.load(file)
+            except json.JSONDecodeError:
+                content = []
+    else:
+        content = []
+
+    # Get the current date and time
+    now = datetime.now()
+    # Format the date and time as a string
+    date_time_string = now.strftime("%Y-%m-%d %H:%M:%S")
+    new_data.append(date_time_string)
+
+    # Append the new data to the content
+    content.append(new_data)
+
+    # Write the updated content back to the file
+    with open(file_path, 'w') as file:
+        json.dump(content, file)
+