Spaces:
Running
Running
File size: 4,675 Bytes
1b870f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import cv2
import numpy as np
import streamlit as st
from run_yolo import get_layout_results
from order_text_blocks import get_ordered_data
from run_ocr import OCR
from main import driver
import json
import pandas as pd
colors = {
'Articles': [0, 0, 0], # Red
'Advertisement': [0, 255, 0], # Green
'Headlines': [0, 0, 255], # Blue
'Sub-headlines': [255, 255, 0], # Yellow
'Graphics': [255, 0, 255], # Magenta
'Images': [128, 0, 128], # Purple
'Tables': [0, 255, 255], # Teal
'Header': [0, 0, 0], # Black
'Text Block': [255, 0, 0]
}
try:
st.set_page_config(layout="wide", page_title="Newspaper Layout Detection and OCR Demo")
st.markdown("<h1 style='text-align: center; color: #333;'>Newspaper Layout Detection and OCR Demo</h1>", unsafe_allow_html=True)
# Streamlit UI for user input
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
language_name = st.text_input("Enter the language name: (hin, en, mal, tel, tam, kan))")
submit_button = st.button("Submit")
# Check if the user clicked the submit button
if submit_button:
# Check if image and language are provided
if uploaded_image is not None and language_name:
# Convert Streamlit file uploader to OpenCV image
image_bytes = uploaded_image.read()
nparr = np.frombuffer(image_bytes, np.uint8)
img_ori = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
img = img_ori.copy()
st.markdown("<p style='text-align: center; color: red'>Image Uploaded Successfully!</p>", unsafe_allow_html=True)
# Continue with the rest of the code...
output_dict, article_ocr_dict = driver(img_ori, language_name, st)
# Create a list to store dictionaries for OCR results
image_data = []
# Visualizing Results
itr = 1
for art_key in article_ocr_dict:
art_coords = art_key.split('_')
art_x1, art_y1, art_x2, art_y2 = int(art_coords[0]), int(art_coords[1]), int(art_coords[2]), int(art_coords[3])
# Mark the article bounding box with dark green color
img_ori = cv2.rectangle(img, (art_x1, art_y1), (art_x2, art_y2), (0, 0, 0), 4)
# Put the article number on the image in large font
img_ori = cv2.putText(img_ori, str(itr), (art_x1, art_y1), cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3, cv2.LINE_AA)
ocr_dict = article_ocr_dict[art_key]
# Initialize variables to store OCR text for each type
headlines_text = ''
subheadlines_text = ''
textblocks_text = ''
for obj_key in ocr_dict:
# obj_key is either of Headlines, Sub-headlines, Text Block
obj_list = ocr_dict[obj_key]
for obj_dict in obj_list:
for key in obj_dict:
coords = key.split('_')
x1, y1, x2, y2 = int(coords[0]), int(coords[1]), int(coords[2]), int(coords[3])
# Mark the bounding box with color corresponding to the object type
img = cv2.rectangle(img, (x1, y1), (x2, y2), colors[obj_key], 2)
# Add the OCR text to the corresponding variable
if obj_key == 'Headlines':
headlines_text += obj_dict[key] + '\n'
elif obj_key == 'Sub-headlines':
subheadlines_text += obj_dict[key] + '\n'
elif obj_key == 'Text Block':
textblocks_text += obj_dict[key] + '\n'
# Add a dictionary to the list for the current article
image_data.append({'Article': itr, 'Headlines': headlines_text, 'Subheadlines': subheadlines_text, 'Textblocks': textblocks_text})
itr += 1
# Create a DataFrame from the list of dictionaries
image_df = pd.DataFrame(image_data)
# Use Streamlit columns to display the image and results side by side
col1, col2 = st.columns(2)
# Display the image with marked bounding boxes in the left column
col1.image(img_ori, use_column_width=True, channels="BGR", caption="Image with Marked Bounding Boxes")
# Display the DataFrame for OCR results for the whole image in the right column
col2.table(image_df.set_index('Article').style.set_table_styles([{'selector': 'th', 'props': [('text-align', 'center')]}]))
except Exception as e:
st.exception(e) |