amit-scans's picture
remove installation of easyocr and transformers through os.system
2993e07
raw
history blame contribute delete
No virus
2.29 kB
import PIL
from PIL import ImageDraw
from PIL import Image
import streamlit as st
import os
def load_image(image_file):
img = PIL.Image.open(image_file)
return img
def init_session_states():
if 'disp' not in st.session_state:
st.session_state['disp'] = st.empty()
st.session_state['disp'].text("Setting up environment with latest build of easyocr. This will take about a minute ")
if 'init' not in st.session_state:
st.session_state['init'] = 1
# Not required as they are already installed through requirements and also seems to cause errors
# os.system('pip install git+git://github.com/jaidedai/easyocr.git')
# os.system('pip install git+https://github.com/huggingface/transformers.git --upgrade')
init_session_states()
import easyocr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
def text_recognition(image):
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
#processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
#model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
pixel_values = processor(image, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
st.write(generated_text)
def main():
st.session_state['disp'].text("Env setup up Complete")
uploaded_file = st.file_uploader("Choose image file to detect text",type=['jpeg','jpg'])
if uploaded_file is not None:
file_details = {"FileName":uploaded_file.name,"FileType":uploaded_file.type,"FileSize":uploaded_file.size}
st.write(file_details)
image = load_image(uploaded_file)
st.image(image,width=500)
st.write("Detecting text bounding box and Take 1 recognition...")
reader = easyocr.Reader(['en'],gpu=True)
bound = reader.readtext(image)
st.write("Bounding box Detection complete")
st.write(str(bound))
st.write("Recognizing text - Take 2....")
text_recognition(image)
if __name__ == "__main__":
main()