TrOCR_demo / app.py
xelpmocAI's picture
Remove Xelpmoc from title
d3d1cec verified
import streamlit as st
from http import client
import os,json
import pandas as pd
import requests
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
st.header("Snippet level OCR")
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-printed')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-printed')
def TrOCR_predict(pixel_values, processor, model):
generated_ids = model.generate(pixel_values,output_scores=True,return_dict_in_generate=True, max_length = 64)
predicted_text = processor.batch_decode(generated_ids[0], skip_special_tokens=True)
return predicted_text
uploaded_file = st.file_uploader("Choose a file")
if uploaded_file is not None:
content = uploaded_file.read()
st.image(uploaded_file)
image = Image.open(uploaded_file)
if image.mode != "RGB": # Convert the image to RGB
image = image.convert("RGB")
pixel_values = processor(images=image, return_tensors="pt").pixel_values
predicted_text = TrOCR_predict(pixel_values, processor, model)[0]
texts = predicted_text
st.write(texts)