Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from ocr_processor import OCRProcessor | |
| import tempfile | |
| import os | |
| from PIL import Image | |
| import json | |
| # Page configuration | |
| st.set_page_config( | |
| page_title="OCR Hub", | |
| page_icon="🔍", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| # Custom CSS for better UI | |
| st.markdown(""" | |
| <style> | |
| .stApp { | |
| max-width: 100%; | |
| padding: 1rem; | |
| } | |
| .main { | |
| background-color: #f8f9fa; | |
| } | |
| .stButton button { | |
| width: 100%; | |
| border-radius: 5px; | |
| height: 3em; | |
| background-color: #4CAF50; | |
| color: white; | |
| } | |
| .stSelectbox { | |
| margin-bottom: 1rem; | |
| } | |
| .upload-text { | |
| text-align: center; | |
| padding: 2rem; | |
| border: 2px dashed #ccc; | |
| border-radius: 10px; | |
| background-color: #ffffff; | |
| } | |
| .stImage { | |
| border-radius: 10px; | |
| box-shadow: 0 2px 4px rgba(0,0,0,0.1); | |
| } | |
| .gallery { | |
| display: grid; | |
| grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); | |
| gap: 1rem; | |
| padding: 1rem; | |
| } | |
| .gallery-item { | |
| border: 1px solid #ddd; | |
| border-radius: 8px; | |
| padding: 0.5rem; | |
| background: white; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| def get_available_models(): | |
| return ["llava:7b", "MiniCPM-V","llama3.2-vision:11b"] | |
| def process_single_image(processor, image_path, format_type, enable_preprocessing): | |
| """Process a single image and return the result""" | |
| try: | |
| result = processor.process_image( | |
| image_path=image_path, | |
| format_type=format_type, | |
| preprocess=enable_preprocessing | |
| ) | |
| return result | |
| except Exception as e: | |
| return f"Error processing image: {str(e)}" | |
| def process_batch_images(processor, image_paths, format_type, enable_preprocessing): | |
| """Process multiple images and return results""" | |
| try: | |
| results = processor.process_batch( | |
| input_path=image_paths, | |
| format_type=format_type, | |
| preprocess=enable_preprocessing | |
| ) | |
| return results | |
| except Exception as e: | |
| return {"error": str(e)} | |
| def main(): | |
| st.title("🔍 OCR Hub") | |
| st.markdown("<p style='text-align: center; color: #666;'>Powered by Ollama Vision Models</p>", unsafe_allow_html=True) | |
| # Sidebar controls | |
| with st.sidebar: | |
| st.header("🎮 Controls") | |
| selected_model = st.selectbox( | |
| "🤖 Select Vision Model", | |
| get_available_models(), | |
| index=0, | |
| ) | |
| format_type = st.selectbox( | |
| "📄 Output Format", | |
| ["markdown", "text", "json", "structured", "key_value"], | |
| help="Choose how you want the extracted text to be formatted" | |
| ) | |
| max_workers = st.slider( | |
| "🔄 Parallel Processing", | |
| min_value=1, | |
| max_value=8, | |
| value=2, | |
| help="Number of images to process in parallel (for batch processing)" | |
| ) | |
| enable_preprocessing = st.checkbox( | |
| "🔍 Enable Preprocessing", | |
| value=True, | |
| help="Apply image enhancement and preprocessing" | |
| ) | |
| st.markdown("---") | |
| # Model info box | |
| if selected_model == "llava:7b": | |
| st.info("LLaVA 7B: Efficient vision-language model optimized for real-time processing") | |
| elif selected_model == "MiniCPM-V": | |
| st.info("MiniCPM-V 2.6: A GPT-4V Level MLLM for Single Image, Multi Image and Video, outperforms GPT-4o mini, Gemini 1.5 Pro and Claude 3.5 Sonnet") | |
| else: | |
| st.info("Llama 3.2 Vision: Advanced model with high accuracy for complex text extraction") | |
| # Initialize OCR Processor | |
| processor = OCRProcessor(model_name=selected_model, max_workers=max_workers) | |
| # Main content area with tabs | |
| tab1, tab2 = st.tabs(["📸 Image Processing", "ℹ️ About"]) | |
| with tab1: | |
| # File upload area with multiple file support | |
| uploaded_files = st.file_uploader( | |
| "Drop your images here", | |
| type=['png', 'jpg', 'jpeg', 'tiff', 'bmp', 'pdf'], | |
| accept_multiple_files=True, | |
| help="Supported formats: PNG, JPG, JPEG, TIFF, BMP, PDF" | |
| ) | |
| if uploaded_files: | |
| # Create a temporary directory for uploaded files | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| image_paths = [] | |
| # Save uploaded files and collect paths | |
| for uploaded_file in uploaded_files: | |
| temp_path = os.path.join(temp_dir, uploaded_file.name) | |
| with open(temp_path, "wb") as f: | |
| f.write(uploaded_file.getvalue()) | |
| image_paths.append(temp_path) | |
| # Display images in a gallery | |
| st.subheader(f"📸 Input Images ({len(uploaded_files)} files)") | |
| cols = st.columns(min(len(uploaded_files), 4)) | |
| for idx, uploaded_file in enumerate(uploaded_files): | |
| with cols[idx % 4]: | |
| image = Image.open(uploaded_file) | |
| st.image(image, use_container_width=True, caption=uploaded_file.name) | |
| # Process button | |
| if st.button("🚀 Process Images"): | |
| with st.spinner("Processing images..."): | |
| if len(image_paths) == 1: | |
| # Single image processing | |
| result = process_single_image( | |
| processor, | |
| image_paths[0], | |
| format_type, | |
| enable_preprocessing | |
| ) | |
| st.subheader("📝 Extracted Text") | |
| st.markdown(result) | |
| # Download button for single result | |
| st.download_button( | |
| "📥 Download Result", | |
| result, | |
| file_name=f"ocr_result.{format_type}", | |
| mime="text/plain" | |
| ) | |
| else: | |
| # Batch processing | |
| results = process_batch_images( | |
| processor, | |
| image_paths, | |
| format_type, | |
| enable_preprocessing | |
| ) | |
| # Display statistics | |
| st.subheader("📊 Processing Statistics") | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| st.metric("Total Images", results['statistics']['total']) | |
| with col2: | |
| st.metric("Successful", results['statistics']['successful']) | |
| with col3: | |
| st.metric("Failed", results['statistics']['failed']) | |
| # Display results | |
| st.subheader("📝 Extracted Text") | |
| for file_path, text in results['results'].items(): | |
| with st.expander(f"Result: {os.path.basename(file_path)}"): | |
| st.markdown(text) | |
| # Display errors if any | |
| if results['errors']: | |
| st.error("⚠️ Some files had errors:") | |
| for file_path, error in results['errors'].items(): | |
| st.warning(f"{os.path.basename(file_path)}: {error}") | |
| # Download all results as JSON | |
| if st.button("📥 Download All Results"): | |
| json_results = json.dumps(results, indent=2) | |
| st.download_button( | |
| "📥 Download Results JSON", | |
| json_results, | |
| file_name="ocr_results.json", | |
| mime="application/json" | |
| ) | |
| with tab2: | |
| st.header("About OCR Hub") | |
| st.markdown(""" | |
| This application uses state-of-the-art vision language models through Ollama to extract text from images. | |
| ### Features: | |
| - 🖼️ Support for multiple image formats | |
| - 📦 Batch processing capability | |
| - 🔄 Parallel processing | |
| - 🔍 Image preprocessing and enhancement | |
| - 📊 Multiple output formats | |
| - 📥 Easy result download | |
| ### Models: | |
| - **LLaVA 7B**: Efficient vision-language model for real-time processing | |
| - **Llama 3.2 Vision**: Advanced model with high accuracy for complex documents | |
| - **MiniCPM-V 2.6**: Process images with any aspect ratio and up to 1.8 million pixels (e.g., 1344x1344) | |
| """) | |
| if __name__ == "__main__": | |
| main() |