Spaces:
Running
Running
| import streamlit as st | |
| from ocr_cpu import extract_text_got, extract_text_qwen, extract_text_llama, clean_extracted_text | |
| import json | |
| # Set up page layout and styling | |
| st.set_page_config(page_title="MultiModel OCR Fusion", layout="centered", page_icon="📄") | |
| st.markdown( | |
| """ | |
| <style> | |
| .reportview-container { background: #f4f4f4; } | |
| .sidebar .sidebar-content { background: #e0e0e0; } | |
| h1 { color: #007BFF; } | |
| .upload-btn { background-color: #007BFF; color: white; padding: 10px; border-radius: 5px; text-align: center; } | |
| </style> | |
| """, unsafe_allow_html=True | |
| ) | |
| # --- Title Section --- | |
| st.title("📄 MultiModel OCR Fusion") | |
| st.write("Upload an image to extract and clean text using multiple OCR models (GOT, Qwen, LLaMA).") | |
| # --- Image Upload Section --- | |
| uploaded_file = st.file_uploader("Upload an image file", type=["jpg", "jpeg", "png"]) | |
| # Model selection | |
| st.sidebar.title("Model Selection") | |
| model_choice = st.sidebar.selectbox("Choose OCR Model", ("GOT", "Qwen", "LLaMA")) | |
| if uploaded_file is not None: | |
| st.image(uploaded_file, caption='Uploaded Image', use_column_width=True) | |
| # Extract text from the image based on selected model | |
| with st.spinner(f"Extracting text using the {model_choice} model..."): | |
| try: | |
| if model_choice == "GOT": | |
| extracted_text = extract_text_got(uploaded_file) | |
| elif model_choice == "Qwen": | |
| extracted_text = extract_text_qwen(uploaded_file) | |
| elif model_choice == "LLaMA": | |
| extracted_text = extract_text_llama(uploaded_file) | |
| # If no text extracted | |
| if not extracted_text.strip(): | |
| st.warning(f"No text extracted using {model_choice}.") | |
| else: | |
| # Clean the extracted text | |
| cleaned_text = clean_extracted_text(extracted_text) | |
| except Exception as e: | |
| st.error(f"Error during text extraction: {str(e)}") | |
| extracted_text, cleaned_text = "", "" | |
| # --- Display Extracted and Cleaned Text --- | |
| st.subheader(f"Extracted Text using {model_choice}") | |
| st.text_area(f"Raw Text ({model_choice})", extracted_text, height=200) | |
| st.subheader("Cleaned Text (AI-processed)") | |
| st.text_area("Cleaned Text", cleaned_text, height=200) | |
| # Save extracted text for further use | |
| if extracted_text: | |
| with open("extracted_text.json", "w") as json_file: | |
| json.dump({"text": extracted_text}, json_file) | |
| # --- Keyword Search --- | |
| st.subheader("Search for Keywords") | |
| keyword = st.text_input("Enter a keyword to search in the extracted text") | |
| if keyword: | |
| if keyword.lower() in cleaned_text.lower(): | |
| st.success(f"Keyword **'{keyword}'** found in the cleaned text!") | |
| else: | |
| st.error(f"Keyword **'{keyword}'** not found.") | |