import spaces
import streamlit as st
from PIL import Image
import torch
from transformers import (
DonutProcessor,
VisionEncoderDecoderModel,
LayoutLMv3Processor,
LayoutLMv3ForSequenceClassification,
AutoProcessor,
AutoModelForCausalLM,
AutoModelForVisualQuestionAnswering
)
from ultralytics import YOLO
import io
import base64
import json
from datetime import datetime
import os
import logging
# Add this near the top of the file, after imports
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
@st.cache_resource
def load_model(model_name):
"""Load the selected model and processor"""
try:
if model_name == "OmniParser":
try:
# Load model directly using official implementation
processor = AutoProcessor.from_pretrained(
"microsoft/OmniParser",
trust_remote_code=True
)
model = AutoModelForVisualQuestionAnswering.from_pretrained(
"microsoft/OmniParser",
trust_remote_code=True,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
if torch.cuda.is_available():
model = model.to("cuda")
st.success("Successfully loaded OmniParser model")
return {
'processor': processor,
'model': model
}
except Exception as e:
st.error(f"Failed to load OmniParser from HuggingFace Hub: {str(e)}")
logger.error(f"OmniParser loading error: {str(e)}", exc_info=True)
return None
elif model_name == "Donut":
processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base")
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base")
# Configure Donut specific parameters
model.config.decoder_start_token_id = processor.tokenizer.bos_token_id
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.vocab_size = len(processor.tokenizer)
return {'model': model, 'processor': processor}
elif model_name == "LayoutLMv3":
processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
model = LayoutLMv3ForSequenceClassification.from_pretrained("microsoft/layoutlmv3-base")
return {'model': model, 'processor': processor}
else:
raise ValueError(f"Unknown model name: {model_name}")
except Exception as e:
st.error(f"Error loading model {model_name}: {str(e)}")
logger.error(f"Error details: {str(e)}", exc_info=True)
return None
@spaces.GPU
@torch.inference_mode()
def analyze_document(image, model_name, models_dict):
"""Analyze document using selected model"""
try:
if models_dict is None:
return {"error": "Model failed to load", "type": "model_error"}
if model_name == "OmniParser":
# Process image with OmniParser
inputs = models_dict['processor'](
images=image,
return_tensors="pt",
)
if torch.cuda.is_available():
inputs = {k: v.to("cuda") if hasattr(v, "to") else v
for k, v in inputs.items()}
# Generate outputs
outputs = models_dict['model'](**inputs)
# Process results
# The exact processing will depend on the model's output format
results = {
"predictions": outputs.logits.softmax(-1).tolist(),
"detected_elements": len(outputs.logits[0]),
"model_output": {
k: v.tolist() if hasattr(v, "tolist") else str(v)
for k, v in outputs.items()
if k != "last_hidden_state" # Skip large tensors
}
}
return results
elif model_name == "Donut":
model = models_dict['model']
processor = models_dict['processor']
# Process image with Donut
pixel_values = processor(image, return_tensors="pt").pixel_values
task_prompt = "<s_cord>analyze the document and extract information</s_cord>"
decoder_input_ids = processor.tokenizer(
task_prompt,
add_special_tokens=False,
return_tensors="pt"
).input_ids
outputs = model.generate(
pixel_values,
decoder_input_ids=decoder_input_ids,
max_length=512,
early_stopping=True,
pad_token_id=processor.tokenizer.pad_token_id,
eos_token_id=processor.tokenizer.eos_token_id,
use_cache=True,
num_beams=4,
bad_words_ids=[[processor.tokenizer.unk_token_id]],
return_dict_in_generate=True
)
sequence = processor.batch_decode(outputs.sequences)[0]
sequence = sequence.replace(task_prompt, "").replace("</s_cord>", "").strip()
try:
result = json.loads(sequence)
except json.JSONDecodeError:
result = {"raw_text": sequence}
return result
elif model_name == "LayoutLMv3":
model = models_dict['model']
processor = models_dict['processor']
# Process image with LayoutLMv3
encoded_inputs = processor(
image,
return_tensors="pt",
add_special_tokens=True,
return_offsets_mapping=True
)
outputs = model(**encoded_inputs)
predictions = outputs.logits.argmax(-1).squeeze().tolist()
# Convert predictions to labels
words = processor.tokenizer.convert_ids_to_tokens(
encoded_inputs.input_ids.squeeze().tolist()
)
result = {
"predictions": [
{
"text": word,
"label": pred
}
for word, pred in zip(words, predictions)
if word not in ["<s>", "</s>", "<pad>"]
],
"confidence_scores": outputs.logits.softmax(-1).max(-1).values.squeeze().tolist()
}
return result
else:
return {"error": f"Unknown model: {model_name}", "type": "model_error"}
except Exception as e:
import traceback
error_details = traceback.format_exc()
logger.error(f"Analysis error: {str(e)}\n{error_details}")
return {
"error": str(e),
"type": "processing_error",
"details": error_details
}
# Set page config with improved layout
st.set_page_config(
page_title="Document Analysis Comparison",
layout="wide",
initial_sidebar_state="expanded"
)
# Add custom CSS for better styling
st.markdown("""
<style>
.stAlert {
margin-top: 1rem;
}
.upload-text {
font-size: 1.2rem;
margin-bottom: 1rem;
}
.model-info {
padding: 1rem;
border-radius: 0.5rem;
background-color: #f8f9fa;
}
</style>
""", unsafe_allow_html=True)
# Title and description
st.title("Document Understanding Model Comparison")
st.markdown("""
Compare different models for document analysis and understanding.
Upload an image and select a model to analyze it.
""")
# Create two columns for layout
col1, col2 = st.columns([1, 1])
with col1:
# File uploader with improved error handling
uploaded_file = st.file_uploader(
"Choose a document image",
type=['png', 'jpg', 'jpeg', 'pdf'],
help="Supported formats: PNG, JPEG, PDF"
)
if uploaded_file is not None:
try:
# Display uploaded image
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Document', use_column_width=True)
except Exception as e:
st.error(f"Error loading image: {str(e)}")
with col2:
# Model selection with detailed information
model_info = {
"Donut": {
"description": "Best for structured OCR and document format understanding",
"memory": "6-8GB",
"strengths": ["Structured OCR", "Memory efficient", "Good with fixed formats"],
"best_for": ["Invoices", "Forms", "Structured documents", "Tables"]
},
"LayoutLMv3": {
"description": "Strong layout understanding with reasoning capabilities",
"memory": "12-15GB",
"strengths": ["Layout understanding", "Reasoning", "Pre-trained knowledge"],
"best_for": ["Complex documents", "Mixed layouts", "Documents with tables", "Multi-column text"]
},
"OmniParser": {
"description": "General screen parsing tool for UI understanding",
"memory": "8-10GB",
"strengths": ["UI element detection", "Interactive element recognition", "Function description"],
"best_for": ["Screenshots", "UI analysis", "Interactive elements", "Web interfaces"]
}
}
selected_model = st.selectbox(
"Select Model",
list(model_info.keys())
)
# Display enhanced model information
st.markdown("### Model Details")
with st.expander("Model Information", expanded=True):
st.markdown(f"**Description:** {model_info[selected_model]['description']}")
st.markdown(f"**Memory Required:** {model_info[selected_model]['memory']}")
st.markdown("**Strengths:**")
for strength in model_info[selected_model]['strengths']:
st.markdown(f"- {strength}")
st.markdown("**Best For:**")
for use_case in model_info[selected_model]['best_for']:
st.markdown(f"- {use_case}")
# Inside the analysis section, replace the existing if-block with:
if uploaded_file is not None and selected_model:
if st.button("Analyze Document", help="Click to start document analysis"):
# Create two columns for results and debug info
result_col, debug_col = st.columns([1, 1])
with st.spinner('Processing...'):
try:
# Create a progress bar in results column
with result_col:
st.markdown("### Analysis Progress")
progress_bar = st.progress(0)
# Initialize debug column
with debug_col:
st.markdown("### Debug Information")
debug_container = st.empty()
def update_debug(message, level="info"):
"""Update debug information with timestamp"""
timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
color = {
"info": "blue",
"warning": "orange",
"error": "red",
"success": "green"
}.get(level, "black")
return f"<div style='color: {color};'>[{timestamp}] {message}</div>"
debug_messages = []
def add_debug(message, level="info"):
debug_messages.append(update_debug(message, level))
debug_container.markdown(
"\n".join(debug_messages),
unsafe_allow_html=True
)
# Load model with progress update
with result_col:
progress_bar.progress(25)
st.info("Loading model...")
add_debug(f"Loading {selected_model} model and processor...")
models_dict = load_model(selected_model)
if models_dict is None:
with result_col:
st.error("Failed to load model. Please try again.")
add_debug("Model loading failed!", "error")
else:
add_debug("Model loaded successfully", "success")
# For device info, we need to check which model we're using
if selected_model == "OmniParser":
model_device = next(models_dict['model'].parameters()).device
else:
model_device = next(models_dict['model'].parameters()).device
add_debug(f"Model device: {model_device}")
# Update progress
with result_col:
progress_bar.progress(50)
st.info("Analyzing document...")
# Log image details
add_debug(f"Image size: {image.size}")
add_debug(f"Image mode: {image.mode}")
# Analyze document
add_debug("Starting document analysis...")
results = analyze_document(image, selected_model, models_dict)
add_debug("Analysis completed", "success")
# Update progress
with result_col:
progress_bar.progress(75)
st.markdown("### Analysis Results")
if isinstance(results, dict) and "error" in results:
st.error(f"Analysis Error: {results['error']}")
add_debug(f"Analysis error: {results['error']}", "error")
else:
# Pretty print the results in results column
st.json(results)
# Show detailed results breakdown in debug column
add_debug("Results breakdown:", "info")
if isinstance(results, dict):
for key, value in results.items():
add_debug(f"- {key}: {type(value)}")
else:
add_debug(f"Result type: {type(results)}")
# Complete progress
progress_bar.progress(100)
st.success("Analysis completed!")
# Final debug info
add_debug("Process completed successfully", "success")
with debug_col:
if torch.cuda.is_available():
st.markdown("### Resource Usage")
st.markdown(f"""
- GPU Memory: {torch.cuda.max_memory_allocated()/1024**2:.2f}MB
- GPU Utilization: {torch.cuda.utilization()}%
""")
except Exception as e:
with result_col:
st.error(f"Error during analysis: {str(e)}")
add_debug(f"Error: {str(e)}", "error")
add_debug(f"Error type: {type(e)}", "error")
if hasattr(e, '__traceback__'):
add_debug("Traceback available in logs", "warning")
# Add improved information about usage and limitations
def verify_weights_directory():
"""Verify the weights directory structure and files"""
weights_path = "weights"
required_files = {
os.path.join(weights_path, "icon_detect", "model.safetensors"): "YOLO model weights",
os.path.join(weights_path, "icon_detect", "model.yaml"): "YOLO model config",
os.path.join(weights_path, "icon_caption_florence", "model.safetensors"): "Florence model weights",
os.path.join(weights_path, "icon_caption_florence", "config.json"): "Florence model config",
os.path.join(weights_path, "icon_caption_florence", "generation_config.json"): "Florence generation config"
}
missing_files = []
for file_path, description in required_files.items():
if not os.path.exists(file_path):
missing_files.append(f"{description} at {file_path}")
if missing_files:
st.warning("Missing required model files:")
for missing in missing_files:
st.write(f"- {missing}")
return False
return True
# Add this in your app's initialization
if st.checkbox("Check Model Files"):
if verify_weights_directory():
st.success("All required model files are present")
else:
st.error("Some model files are missing. Please ensure all required files are in the weights directory")
st.markdown("""
---
### Usage Notes:
- Different models excel at different types of documents
- Processing time and memory requirements vary by model
- Image quality significantly affects results
- Some models may require specific document formats
""")
# Add performance metrics section
if st.checkbox("Show Performance Metrics"):
st.markdown("""
### Model Performance Metrics
| Model | Avg. Processing Time | Memory Usage | Accuracy* |
|-------|---------------------|--------------|-----------|
| Donut | 2-3 seconds | 6-8GB | 85-90% |
| LayoutLMv3 | 3-4 seconds | 12-15GB | 88-93% |
| OmniParser | 2-3 seconds | 8-10GB | 85-90% |
*Accuracy varies based on document type and quality
""")
# Add a footer with version and contact information
st.markdown("---")
st.markdown("""
v1.1 - Created with Streamlit
\nPowered by Hugging Face Spaces 🤗
""")
# Add model selection guidance
if st.checkbox("Show Model Selection Guide"):
st.markdown("""
### How to Choose the Right Model
1. **Donut**: Choose for structured documents with clear layouts
2. **LayoutLMv3**: Best for documents with complex layouts and relationships
3. **OmniParser**: Best for UI elements and screen parsing
""")