|
import pathlib |
|
import streamlit as st |
|
from PIL import Image, ImageDraw, ImageFont |
|
from streamlit_image_coordinates import streamlit_image_coordinates |
|
import pytesseract |
|
import pypdfium2 |
|
|
|
from streamlit.runtime.uploaded_file_manager import UploadedFile |
|
|
|
|
|
LANG = "eng+por" |
|
OPTIONS = "" |
|
FONT_FILE = pathlib.Path(__file__).parent / "Roboto-Regular.ttf" |
|
|
|
st.set_page_config( |
|
page_title="Streamlit Demo", |
|
layout="wide", |
|
page_icon="🪟", |
|
) |
|
|
|
"# Streamlit PDF OCR Demo" |
|
|
|
@st.cache_data(hash_funcs={"file": lambda uploaded_file: hash(uploaded_file.file_id)}) |
|
def convert_pdf_to_img(file: UploadedFile) -> Image.Image: |
|
doc = pypdfium2.PdfDocument(file) |
|
raw_image: pypdfium2.PdfBitmap = doc[0].render(scale=2, grayscale=True) |
|
image: Image.Image = raw_image.to_pil() |
|
return image |
|
|
|
|
|
@st.cache_data |
|
def load_image_and_ocr(img: Image.Image, lang: str, config: str) -> tuple[Image.Image, dict, list[dict]]: |
|
data = pytesseract.image_to_data(img, lang, config, output_type=pytesseract.Output.DICT) |
|
_boxes = [{"text": text} for text in data["text"]] |
|
for key in ("left", "top", "width", "height"): |
|
for i, val in enumerate(data[key]): |
|
_boxes[i][key] = val |
|
boxes = [box for box in _boxes if box['text'].strip() != ''] |
|
return img.convert("RGB"), data, boxes |
|
|
|
|
|
def main(image: Image.Image): |
|
img, data, boxes = load_image_and_ocr(image, LANG, OPTIONS) |
|
|
|
if (coords := st.session_state.get("image_coords")) is None: |
|
coords = {"x": 0, "y": 0} |
|
|
|
draw = ImageDraw.Draw(img) |
|
text = None |
|
for box in boxes: |
|
MARGIN = 5 |
|
x1, y1, width, height = (box['left'], box['top'], box['width'], box['height']) |
|
x2, y2 = x1+width, y1+height |
|
x1, y1, x2, y2 = (x1-MARGIN, y1-MARGIN, x2+MARGIN, y2+MARGIN) |
|
|
|
if (x1 <= coords["x"] <= x2) and (y1 <= coords["y"] <= y2): |
|
color = "blue" |
|
font = ImageFont.FreeTypeFont(FONT_FILE, size=24) |
|
text = box['text'] |
|
|
|
draw.text(((x1+x2)/2, y1-15), text, align="center", anchor="mm", fill="red", font=font) |
|
else: |
|
color = "green" |
|
|
|
draw.rectangle((x1, y1, x2, y2), fill=None, outline=color, width=2) |
|
|
|
"## Click at a rectangle" |
|
if text: |
|
f'''### Selected text: "{text}"''' |
|
else: |
|
'''### Selected text will appear here''' |
|
|
|
|
|
streamlit_image_coordinates(img, key="image_coords", click_and_drag=False) |
|
|
|
st.file_uploader("Upload PDF", ".pdf", key="pdf") |
|
|
|
if (pdf_file := st.session_state.get("pdf")) is not None: |
|
base_image = convert_pdf_to_img(pdf_file) |
|
main(base_image) |
|
|
