myocr / src /app.py
robbyzhaox's picture
fix rectangle
1f36a7b unverified
import logging
import streamlit as st
import pandas as pd
import io
from PIL import Image, ImageDraw, ImageFont
from myocr.pipelines import CommonOCRPipeline
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)
st.set_page_config(layout="wide")
st.title("MyOCR Demo")
# init pipeline
@st.cache_resource
def load_pipeline():
return CommonOCRPipeline("cpu")
pipeline = load_pipeline()
font = ImageFont.truetype("src/NotoSans.ttf", 12)
def process_image(image: Image.Image, format):
buffer = io.BytesIO()
image.save(buffer, format=format)
results = pipeline(buffer.getvalue())
logger.info(f"ocr results: {results}")
if not results or not results.regions:
return None, []
image_with_boxes = image.copy()
draw = ImageDraw.Draw(image_with_boxes)
table_data = []
for item in results.regions:
if item.confidence <= 0.5:
continue
shape = item.bounding_shape
points = shape.points
if shape.type == "rectangle":
draw.rectangle([(points[0].x, points[0].y), (points[2].x, points[2].y)], outline="red", width=2)
text_pos_y = max(points[0].y - 18, 0)
draw.text((points[0].x, text_pos_y), item.text, font=font, fill="green")
table_data.append((item.text, item.confidence))
return image_with_boxes, table_data
def main():
left_col, right_col = st.columns([2, 1.5])
with left_col:
uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
if not uploaded_file:
return
try:
image = Image.open(uploaded_file).convert("RGB")
except Exception:
st.error("⚠️ Invalid image file. Please upload a valid PNG or JPG.")
logger.warning("Invalid file uploaded.")
return
spinner_container = st.empty()
image_slot = st.empty()
image_slot.image(image, use_container_width=True)
with spinner_container:
with st.spinner("Recognizing text..."):
mime_type = uploaded_file.type
processed_image, table_data = process_image(image, mime_type.split("/")[-1])
if not table_data:
right_col.warning("No text detected.")
return
image_slot.image(processed_image, use_container_width=True)
# --- Show results ---
with right_col:
st.subheader("Recognized Text")
df = pd.DataFrame(table_data, columns=["Text", "Confidence"])
st.table(df)
if __name__ == "__main__":
try:
main()
except Exception as e:
logger.exception("Unhandled exception occurred.")
st.error(f"Internal Error!")