optimised-ocr / app.py
AuditEdge's picture
doc upload option added
f8afc9b
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Dict
import os
import shutil
import logging
import torch
from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
from dotenv import load_dotenv
import os
from utils import doc_processing
# Load .env file
load_dotenv()
# Access variables
dummy_key = os.getenv("dummy_key")
HUGGINGFACE_AUTH_TOKEN = dummy_key
# Hugging Face model and token
aadhar_model = "AuditEdge/doc_ocr_a" # Replace with your fine-tuned model if applicable
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load the processor (tokenizer + image processor)
processor_aadhar = LayoutLMv3Processor.from_pretrained(
aadhar_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
aadhar_model = LayoutLMv3ForTokenClassification.from_pretrained(
aadhar_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
aadhar_model = aadhar_model.to(device)
# pan model
pan_model = "AuditEdge/doc_ocr_p" # Replace with your fine-tuned model if applicable
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load the processor (tokenizer + image processor)
processor_pan = LayoutLMv3Processor.from_pretrained(
pan_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
pan_model = LayoutLMv3ForTokenClassification.from_pretrained(
pan_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
pan_model = pan_model.to(device)
#
# gst model
gst_model = "AuditEdge/doc_ocr_new_g" # Replace with your fine-tuned model if applicable
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load the processor (tokenizer + image processor)
processor_gst = LayoutLMv3Processor.from_pretrained(
gst_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
gst_model = LayoutLMv3ForTokenClassification.from_pretrained(
gst_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
gst_model = gst_model.to(device)
#cheque model
cheque_model = "AuditEdge/doc_ocr_new_c" # Replace with your fine-tuned model if applicable
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load the processor (tokenizer + image processor)
processor_cheque = LayoutLMv3Processor.from_pretrained(
cheque_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
cheque_model = LayoutLMv3ForTokenClassification.from_pretrained(
cheque_model,
use_auth_token=HUGGINGFACE_AUTH_TOKEN
)
cheque_model = cheque_model.to(device)
# Verify model and processor are loaded
print("Model and processor loaded successfully!")
print(f"Model is on device: {next(aadhar_model.parameters()).device}")
# Import inference modules
from layoutlmv3FineTuning.Layoutlm_inference.ocr import prepare_batch_for_inference
from layoutlmv3FineTuning.Layoutlm_inference.inference_handler import handle
# Create FastAPI instance
app = FastAPI(debug=True)
# Enable CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Configure directories
UPLOAD_FOLDER = './uploads/'
processing_folder = "./processed_images"
os.makedirs(UPLOAD_FOLDER, exist_ok=True) # Ensure the main upload folder exists
os.makedirs(processing_folder,exist_ok=True)
UPLOAD_DIRS = {
"aadhar_file": "uploads/aadhar/",
"pan_file": "uploads/pan/",
"cheque_file": "uploads/cheque/",
"gst_file": "uploads/gst/",
}
process_dirs = {
"aadhar_file": "processed_images/aadhar/",
"pan_file": "processed_images/pan/",
"cheque_file": "processed_images/cheque/",
"gst_file": "processed_images/gst/",
}
# Ensure individual directories exist
for dir_path in UPLOAD_DIRS.values():
os.makedirs(dir_path, exist_ok=True)
for dir_path in process_dirs.values():
os.makedirs(dir_path, exist_ok=True)
# Logger configuration
logging.basicConfig(level=logging.INFO)
# Perform Inference
def perform_inference(file_paths: Dict[str, str]):
# Dictionary to map document types to their respective model directories
model_dirs = {
"aadhar_file": aadhar_model,
"pan_file": pan_model,
"cheque_file": cheque_model,
"gst_file": gst_model,
}
# Dictionary to store results for each document type
inference_results = {}
# Loop through the file paths and perform inference
for doc_type, file_path in file_paths.items():
if doc_type in model_dirs:
print(f"Processing {doc_type} using model at {model_dirs[doc_type]}")
# Prepare batch for inference
images_path = [file_path]
inference_batch = prepare_batch_for_inference(images_path)
# Prepare context for the specific document type
# context = {"model_dir": model_dirs[doc_type]}
# context = aadhar_model
if doc_type == "aadhar_file":
context = aadhar_model
processor = processor_aadhar
name = "aadhar"
attachemnt_num = 3
if doc_type == "pan_file":
context = pan_model
processor = processor_pan
name = "pan"
attachemnt_num = 2
if doc_type == "gst_file":
context = gst_model
processor = processor_gst
name = "gst"
attachemnt_num = 4
if doc_type == "cheque_file":
context = cheque_model
processor = processor_cheque
name = "cheque"
attachemnt_num = 8
# Perform inference (replace `handle` with your actual function)
result = handle(inference_batch, context,processor,name)
# Store the result
inference_results["attachment_{}".format(attachemnt_num)] = result
else:
print(f"Model directory not found for {doc_type}. Skipping.")
return inference_results
# Routes
@app.get("/")
def greet_json():
return {"Hello": "World!"}
@app.post("/api/aadhar_ocr")
async def aadhar_ocr(
aadhar_file: UploadFile = File(None),
pan_file: UploadFile = File(None),
cheque_file: UploadFile = File(None),
gst_file: UploadFile = File(None),
):
try:
# Handle file uploads
file_paths = {}
for file_type, folder in UPLOAD_DIRS.items():
file = locals()[file_type] # Dynamically access the file arguments
if file:
# Save the file in the respective directory
file_path = os.path.join(folder, file.filename)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
file_paths[file_type] = file_path
# Log received files
logging.info(f"Received files: {list(file_paths.keys())}")
print("file_paths",file_paths)
files = {}
for key, value in file_paths.items():
name = value.split("/")[-1].split(".")[0]
id_type = key.split("_")[0]
doc_type = value.split("/")[-1].split(".")[1]
f_path = value
preprocessing = doc_processing(name,id_type,doc_type,f_path)
response = preprocessing.process()
files[key] = response["output_p"]
print("response",response)
# Perform inference
result = perform_inference(files)
return {"status": "success", "result": result}
except Exception as e:
logging.error(f"Error processing files: {e}")
# raise HTTPException(status_code=500, detail="Internal Server Error")
return {"status":400}