Spaces:
Runtime error
Runtime error
| 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) |