OCR / deploy.py
ShahzainHaider's picture
Upload folder using huggingface_hub
01162ca
raw
history blame
2.87 kB
import numpy as np
import gradio as gr
from fastapi import FastAPI, File, UploadFile
from app.extract_country.country_dictionary import find_country
from app.extract_dates.date_engine import date_extractor
from app.extract_gender.gender_extractor import gender_extract
from app.extract_identity_number.doc_number_extractor import doc_number
from app.ocr_engine.ocr import OCR
from paddleocr import PaddleOCR
from app.layoutLM_api.api import custom_ocr
import cv2
def ocr(image_input, webcam_input):
file_path = 'app/images/idcards/input.jpg'
print("image_input : ", type(image_input))
print("Webcam Image : ", type(webcam_input))
if image_input is not None and image_input.any():
print("In input image")
cv2.imwrite(file_path, image_input)
else:
print("In webcam")
cv2.imwrite(file_path, webcam_input)
# print(type(input_img))
ocr = PaddleOCR(use_angle_cls=True, lang='en') # need to run only once to download and load model into memory
# with open(file_path, "wb+") as file_object:
# file_object.write(input_img.file.read())
dictionary = custom_ocr(file_path)
if '' in list(dictionary.values()):
print("Missing value found in Dic")
extract_text = []
result = ocr.ocr(file_path, cls=True)
extract_text = [line[1][0] for res in result for line in res]
print("extract_text", extract_text)
if len(dictionary['gender']) == 0:
print("Gender Missing")
gender_found, gender, ocr_list = gender_extract(extract_text)
dictionary["gender"] = gender if gender_found else None
if len(dictionary['dob']) == 0:
print("Dob Missing")
dob_found, dob, ocr_list = date_extractor(extract_text)
dictionary["dob"] = dob if dob_found else None
if len(dictionary['country']) == 0:
print("Country Missing")
country_found, country, ocr_list = find_country(extract_text)
dictionary["country"] = country if country_found else None
if len(dictionary['document_number']) == 0:
print("document Number missing")
document_number_found, document_number, ocr_list = doc_number(extract_text)
dictionary["document_number"] = document_number if document_number_found else None
print("Updated Dict ",dictionary)
response = {"Status" : 200, "OCR" : dictionary}
return response
# webcam_checkbox = gr.inputs.Checkbox(label="Use webcam", optional=True)
# Define the input objects
image_input = gr.inputs.Image(label="Upload Image")
webcam_input = gr.inputs.Image(label="Webcam", source="webcam")
# Create the Gradio interface
interface = gr.Interface(fn=ocr, inputs=[image_input, webcam_input], outputs="json")
interface.launch(share=True)