#!/usr/bin/env python3
"""
Dots.OCR Gradio Demo Application
A Gradio-based web interface for demonstrating the Dots.OCR model using Hugging Face transformers.
This application provides OCR and layout analysis capabilities for documents and images.
"""
import os
import json
import traceback
import math
from io import BytesIO
from typing import Optional, Dict, Any, Tuple, List
import requests
# Set LOCAL_RANK for transformers
if "LOCAL_RANK" not in os.environ:
os.environ["LOCAL_RANK"] = "0"
import torch
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoModelForCausalLM, AutoProcessor
from qwen_vl_utils import process_vision_info
import fitz # PyMuPDF
# Constants
MIN_PIXELS = 3136
MAX_PIXELS = 11289600
IMAGE_FACTOR = 28
# Prompts
dict_promptmode_to_prompt = {
"prompt_layout_all_en": """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
1. Bbox format: [x1, y1, x2, y2]
2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
3. Text Extraction & Formatting Rules:
- Picture: For the 'Picture' category, the text field should be omitted.
- Formula: Format its text as LaTeX.
- Table: Format its text as HTML.
- All Others (Text, Title, etc.): Format their text as Markdown.
4. Constraints:
- The output text must be the original text from the image, with no translation.
- All layout elements must be sorted according to human reading order.
5. Final Output: The entire output must be a single JSON object.
""",
"prompt_layout_only_en": """Please output the layout information from this PDF image, including each layout's bbox and its category. The bbox should be in the format [x1, y1, x2, y2]. The layout categories for the PDF document include ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']. Do not output the corresponding text. The layout result should be in JSON format.""",
"prompt_ocr": """Extract the text content from this image.""",
"prompt_grounding_ocr": """Extract text from the given bounding box on the image (format: [x1, y1, x2, y2]).\nBounding Box:\n""",
}
# Utility functions
def round_by_factor(number: int, factor: int) -> int:
"""Returns the closest integer to 'number' that is divisible by 'factor'."""
return round(number / factor) * factor
def smart_resize(
height: int,
width: int,
factor: int = 28,
min_pixels: int = 3136,
max_pixels: int = 11289600,
):
"""Rescales the image so that the following conditions are met:
1. Both dimensions (height and width) are divisible by 'factor'.
2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
3. The aspect ratio of the image is maintained as closely as possible.
"""
if max(height, width) / min(height, width) > 200:
raise ValueError(
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
)
h_bar = max(factor, round_by_factor(height, factor))
w_bar = max(factor, round_by_factor(width, factor))
if h_bar * w_bar > max_pixels:
beta = math.sqrt((height * width) / max_pixels)
h_bar = round_by_factor(height / beta, factor)
w_bar = round_by_factor(width / beta, factor)
elif h_bar * w_bar < min_pixels:
beta = math.sqrt(min_pixels / (height * width))
h_bar = round_by_factor(height * beta, factor)
w_bar = round_by_factor(width * beta, factor)
return h_bar, w_bar
def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
"""Fetch and process an image"""
if isinstance(image_input, str):
if image_input.startswith(("http://", "https://")):
response = requests.get(image_input)
image = Image.open(BytesIO(response.content)).convert('RGB')
else:
image = Image.open(image_input).convert('RGB')
elif isinstance(image_input, Image.Image):
image = image_input.convert('RGB')
else:
raise ValueError(f"Invalid image input type: {type(image_input)}")
if min_pixels is not None or max_pixels is not None:
min_pixels = min_pixels or MIN_PIXELS
max_pixels = max_pixels or MAX_PIXELS
height, width = smart_resize(
image.height,
image.width,
factor=IMAGE_FACTOR,
min_pixels=min_pixels,
max_pixels=max_pixels
)
image = image.resize((width, height), Image.LANCZOS)
return image
def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
"""Load images from PDF file"""
images = []
try:
pdf_document = fitz.open(pdf_path)
for page_num in range(len(pdf_document)):
page = pdf_document.load_page(page_num)
# Convert page to image
mat = fitz.Matrix(2.0, 2.0) # Increase resolution
pix = page.get_pixmap(matrix=mat)
img_data = pix.tobytes("ppm")
image = Image.open(BytesIO(img_data)).convert('RGB')
images.append(image)
pdf_document.close()
except Exception as e:
print(f"Error loading PDF: {e}")
return []
return images
def draw_layout_on_image(image: Image.Image, layout_data: List[Dict]) -> Image.Image:
"""Draw layout bounding boxes on image"""
img_copy = image.copy()
draw = ImageDraw.Draw(img_copy)
# Colors for different categories
colors = {
'Caption': '#FF6B6B',
'Footnote': '#4ECDC4',
'Formula': '#45B7D1',
'List-item': '#96CEB4',
'Page-footer': '#FFEAA7',
'Page-header': '#DDA0DD',
'Picture': '#FFD93D',
'Section-header': '#6C5CE7',
'Table': '#FD79A8',
'Text': '#74B9FF',
'Title': '#E17055'
}
try:
# Load a font
try:
font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
except Exception:
font = ImageFont.load_default()
for item in layout_data:
if 'bbox' in item and 'category' in item:
bbox = item['bbox']
category = item['category']
color = colors.get(category, '#000000')
# Draw rectangle
draw.rectangle(bbox, outline=color, width=2)
# Draw label
label = category
label_bbox = draw.textbbox((0, 0), label, font=font)
label_width = label_bbox[2] - label_bbox[0]
label_height = label_bbox[3] - label_bbox[1]
# Position label above the box
label_x = bbox[0]
label_y = max(0, bbox[1] - label_height - 2)
# Draw background for label
draw.rectangle(
[label_x, label_y, label_x + label_width + 4, label_y + label_height + 2],
fill=color
)
# Draw text
draw.text((label_x + 2, label_y + 1), label, fill='white', font=font)
except Exception as e:
print(f"Error drawing layout: {e}")
return img_copy
def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text', no_page_hf: bool = False) -> str:
"""Convert layout JSON to markdown format"""
markdown_lines = []
if not no_page_hf:
markdown_lines.append("# Document Content\n")
try:
# Sort items by reading order (top to bottom, left to right)
sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
for item in sorted_items:
category = item.get('category', '')
text = item.get(text_key, '')
if not text:
continue
if category == 'Title':
markdown_lines.append(f"# {text}\n")
elif category == 'Section-header':
markdown_lines.append(f"## {text}\n")
elif category == 'Text':
markdown_lines.append(f"{text}\n")
elif category == 'List-item':
markdown_lines.append(f"- {text}\n")
elif category == 'Table':
# If text is already HTML, keep it as is
if text.strip().startswith('<'):
markdown_lines.append(f"{text}\n")
else:
markdown_lines.append(f"**Table:** {text}\n")
elif category == 'Formula':
# If text is LaTeX, format it properly
if text.strip().startswith('$') or '\\' in text:
markdown_lines.append(f"$$\n{text}\n$$\n")
else:
markdown_lines.append(f"**Formula:** {text}\n")
elif category == 'Caption':
markdown_lines.append(f"*{text}*\n")
elif category == 'Footnote':
markdown_lines.append(f"^{text}^\n")
elif category in ['Page-header', 'Page-footer']:
# Skip headers and footers in main content
continue
else:
markdown_lines.append(f"{text}\n")
markdown_lines.append("") # Add spacing
except Exception as e:
print(f"Error converting to markdown: {e}")
return str(layout_data)
return "\n".join(markdown_lines)
# Initialize model and processor at script level
model_id = "rednote-hilab/dots.ocr"
model = AutoModelForCausalLM.from_pretrained(
model_id,
attn_implementation="flash_attention_2",
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True
)
processor = AutoProcessor.from_pretrained(
model_id,
trust_remote_code=True
)
# Global state variables
device = "cuda" if torch.cuda.is_available() else "cpu"
# PDF handling state
pdf_cache = {
"images": [],
"current_page": 0,
"total_pages": 0,
"file_type": None,
"is_parsed": False,
"results": []
}
# Processing state
processing_results = {
'original_image': None,
'processed_image': None,
'layout_result': None,
'markdown_content': None,
'raw_output': None,
}
def inference(image: Image.Image, prompt: str, max_new_tokens: int = 24000) -> str:
"""Run inference on an image with the given prompt"""
try:
if model is None or processor is None:
raise RuntimeError("Model not loaded. Please check model initialization.")
# Prepare messages in the expected format
messages = [
{
"role": "user",
"content": [
{
"type": "image",
"image": image
},
{"type": "text", "text": prompt}
]
}
]
# Apply chat template
text = processor.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
# Process vision information
image_inputs, video_inputs = process_vision_info(messages)
# Prepare inputs
inputs = processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
)
# Move to device
inputs = inputs.to(device)
# Generate output
with torch.no_grad():
generated_ids = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=False,
temperature=0.1
)
# Decode output
generated_ids_trimmed = [
out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
generated_ids_trimmed,
skip_special_tokens=True,
clean_up_tokenization_spaces=False
)
return output_text[0] if output_text else ""
except Exception as e:
print(f"Error during inference: {e}")
traceback.print_exc()
return f"Error during inference: {str(e)}"
def process_image(
image: Image.Image,
prompt_mode: str,
min_pixels: Optional[int] = None,
max_pixels: Optional[int] = None
) -> Dict[str, Any]:
"""Process a single image with the specified prompt mode"""
try:
# Resize image if needed
if min_pixels is not None or max_pixels is not None:
image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
# Get prompt
prompt = dict_promptmode_to_prompt[prompt_mode]
# Run inference
raw_output = inference(image, prompt)
# Process results based on prompt mode
result = {
'original_image': image,
'raw_output': raw_output,
'prompt_mode': prompt_mode,
'processed_image': image,
'layout_result': None,
'markdown_content': None
}
# For layout analysis prompts, try to parse JSON and create visualizations
if prompt_mode in ['prompt_layout_all_en', 'prompt_layout_only_en']:
try:
# Try to parse JSON output
layout_data = json.loads(raw_output)
result['layout_result'] = layout_data
# Create visualization with bounding boxes
try:
processed_image = draw_layout_on_image(image, layout_data)
result['processed_image'] = processed_image
except Exception as e:
print(f"Error drawing layout: {e}")
result['processed_image'] = image
# Generate markdown if text is available
if prompt_mode == 'prompt_layout_all_en':
try:
markdown_content = layoutjson2md(image, layout_data, text_key='text')
result['markdown_content'] = markdown_content
except Exception as e:
print(f"Error generating markdown: {e}")
result['markdown_content'] = raw_output
except json.JSONDecodeError:
print("Failed to parse JSON output, using raw output")
result['markdown_content'] = raw_output
else:
# For OCR prompts, use raw output as markdown
result['markdown_content'] = raw_output
return result
except Exception as e:
print(f"Error processing image: {e}")
traceback.print_exc()
return {
'original_image': image,
'raw_output': f"Error processing image: {str(e)}",
'prompt_mode': prompt_mode,
'processed_image': image,
'layout_result': None,
'markdown_content': f"Error processing image: {str(e)}"
}
def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
"""Load file for preview (supports PDF and images)"""
global pdf_cache
if not file_path or not os.path.exists(file_path):
return None, "No file selected"
file_ext = os.path.splitext(file_path)[1].lower()
try:
if file_ext == '.pdf':
# Load PDF pages
images = load_images_from_pdf(file_path)
if not images:
return None, "Failed to load PDF"
pdf_cache.update({
"images": images,
"current_page": 0,
"total_pages": len(images),
"file_type": "pdf",
"is_parsed": False,
"results": []
})
return images[0], f"Page 1 / {len(images)}"
elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
# Load single image
image = Image.open(file_path).convert('RGB')
pdf_cache.update({
"images": [image],
"current_page": 0,
"total_pages": 1,
"file_type": "image",
"is_parsed": False,
"results": []
})
return image, "Page 1 / 1"
else:
return None, f"Unsupported file format: {file_ext}"
except Exception as e:
print(f"Error loading file: {e}")
return None, f"Error loading file: {str(e)}"
def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, str]:
"""Navigate through PDF pages"""
global pdf_cache
if not pdf_cache["images"]:
return None, "No file loaded", "No results yet"
if direction == "prev":
pdf_cache["current_page"] = max(0, pdf_cache["current_page"] - 1)
elif direction == "next":
pdf_cache["current_page"] = min(
pdf_cache["total_pages"] - 1,
pdf_cache["current_page"] + 1
)
index = pdf_cache["current_page"]
current_image = pdf_cache["images"][index]
page_info = f"Page {index + 1} / {pdf_cache['total_pages']}"
# Get results for current page if available
current_result = ""
if (pdf_cache["is_parsed"] and
index < len(pdf_cache["results"]) and
pdf_cache["results"][index]):
result = pdf_cache["results"][index]
if result.get('markdown_content'):
current_result = result['markdown_content']
else:
current_result = result.get('raw_output', 'No content available')
else:
current_result = "Page not processed yet"
return current_image, page_info, current_result
def create_gradio_interface():
"""Create the Gradio interface"""
# Custom CSS
css = """
.main-container {
max-width: 1400px;
margin: 0 auto;
}
.header-text {
text-align: center;
color: #2c3e50;
margin-bottom: 20px;
}
.process-button {
background: linear-gradient(45deg, #667eea 0%, #764ba2 100%) !important;
border: none !important;
color: white !important;
font-weight: bold !important;
}
.process-button:hover {
transform: translateY(-2px) !important;
box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
}
.info-box {
background: #f8f9fa;
border: 1px solid #dee2e6;
border-radius: 8px;
padding: 15px;
margin: 10px 0;
}
.page-info {
text-align: center;
padding: 8px 16px;
background: #e9ecef;
border-radius: 20px;
font-weight: bold;
margin: 10px 0;
}
.model-status {
padding: 10px;
border-radius: 8px;
margin: 10px 0;
text-align: center;
font-weight: bold;
}
.status-loading {
background: #fff3cd;
color: #856404;
border: 1px solid #ffeaa7;
}
.status-ready {
background: #d1edff;
color: #0c5460;
border: 1px solid #b8daff;
}
.status-error {
background: #f8d7da;
color: #721c24;
border: 1px solid #f5c6cb;
}
"""
with gr.Blocks(theme=gr.themes.Soft(), css=css, title="Dots.OCR Demo") as demo:
# Header
gr.HTML("""
""")
# Model status
model_status = gr.HTML(
'🔄 Initializing model...
',
elem_id="model_status"
)
# Main interface
with gr.Row():
# Left column - Input and controls
with gr.Column(scale=1):
gr.Markdown("### 📁 Input")
# File input
file_input = gr.File(
label="Upload Image or PDF",
file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"],
type="filepath"
)
# Image preview
image_preview = gr.Image(
label="Preview",
type="pil",
interactive=False,
height=300
)
# Page navigation for PDFs
with gr.Row():
prev_page_btn = gr.Button("◀ Previous", size="sm")
page_info = gr.HTML('No file loaded
')
next_page_btn = gr.Button("Next ▶", size="sm")
gr.Markdown("### ⚙️ Settings")
# Prompt mode selection
prompt_mode = gr.Dropdown(
choices=list(dict_promptmode_to_prompt.keys()),
value="prompt_layout_all_en",
label="Task Mode",
info="Choose the type of analysis to perform"
)
# Advanced settings
with gr.Accordion("Advanced Settings", open=False):
max_new_tokens = gr.Slider(
minimum=1000,
maximum=32000,
value=24000,
step=1000,
label="Max New Tokens",
info="Maximum number of tokens to generate"
)
min_pixels = gr.Number(
value=MIN_PIXELS,
label="Min Pixels",
info="Minimum image resolution"
)
max_pixels = gr.Number(
value=MAX_PIXELS,
label="Max Pixels",
info="Maximum image resolution"
)
# Process button
process_btn = gr.Button(
"🚀 Process Document",
variant="primary",
elem_classes=["process-button"],
size="lg"
)
# Clear button
clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
# Right column - Results
with gr.Column(scale=2):
gr.Markdown("### 📊 Results")
# Results tabs
with gr.Tabs():
# Processed image tab
with gr.Tab("🖼️ Processed Image"):
processed_image = gr.Image(
label="Image with Layout Detection",
type="pil",
interactive=False,
height=500
)
# Markdown output tab
with gr.Tab("📝 Extracted Content"):
markdown_output = gr.Markdown(
value="Click 'Process Document' to see extracted content...",
height=500
)
# Raw output tab
with gr.Tab("🔧 Raw Output"):
raw_output = gr.Textbox(
label="Raw Model Output",
lines=20,
max_lines=30,
interactive=False
)
# JSON layout tab
with gr.Tab("📋 Layout JSON"):
json_output = gr.JSON(
label="Layout Analysis Results",
value=None
)
# Prompt display
gr.Markdown("### 💬 Current Prompt")
prompt_display = gr.Textbox(
value=dict_promptmode_to_prompt["prompt_layout_all_en"],
label="Prompt Text",
lines=8,
interactive=False,
info="This is the prompt that will be sent to the model"
)
# Event handlers
def load_model_on_startup():
"""Load model when the interface starts"""
try:
# Model is already loaded at script level
return '✅ Model loaded successfully!
'
except Exception as e:
return f'❌ Error: {str(e)}
'
def process_document(file_path, prompt_mode_val, max_tokens, min_pix, max_pix):
"""Process the uploaded document"""
global pdf_cache
try:
if not file_path:
return (
None,
"Please upload a file first.",
"No file uploaded",
None,
'❌ No file uploaded
'
)
if model is None:
return (
None,
"Model not loaded. Please refresh the page and try again.",
"Model not loaded",
None,
'❌ Model not loaded
'
)
# Load and preview file
image, page_info = load_file_for_preview(file_path)
if image is None:
return (
None,
page_info,
"Failed to load file",
None,
'❌ Failed to load file
'
)
# Process the image(s)
if pdf_cache["file_type"] == "pdf":
# Process all pages for PDF
all_results = []
all_markdown = []
for i, img in enumerate(pdf_cache["images"]):
result = process_image(
img,
prompt_mode_val,
min_pixels=int(min_pix) if min_pix else None,
max_pixels=int(max_pix) if max_pix else None
)
all_results.append(result)
if result.get('markdown_content'):
all_markdown.append(f"## Page {i+1}\n\n{result['markdown_content']}")
pdf_cache["results"] = all_results
pdf_cache["is_parsed"] = True
# Show results for first page
first_result = all_results[0]
combined_markdown = "\n\n---\n\n".join(all_markdown)
return (
first_result['processed_image'],
combined_markdown,
first_result['raw_output'],
first_result['layout_result'],
'✅ Processing completed!
'
)
else:
# Process single image
result = process_image(
image,
prompt_mode_val,
min_pixels=int(min_pix) if min_pix else None,
max_pixels=int(max_pix) if max_pix else None
)
pdf_cache["results"] = [result]
pdf_cache["is_parsed"] = True
return (
result['processed_image'],
result['markdown_content'] or "No content extracted",
result['raw_output'],
result['layout_result'],
'✅ Processing completed!
'
)
except Exception as e:
error_msg = f"Error processing document: {str(e)}"
print(error_msg)
traceback.print_exc()
return (
None,
error_msg,
error_msg,
None,
f'❌ {error_msg}
'
)
def update_prompt_display(mode):
"""Update the prompt display when mode changes"""
return dict_promptmode_to_prompt[mode]
def handle_file_upload(file_path):
"""Handle file upload and show preview"""
if not file_path:
return None, "No file loaded"
image, page_info = load_file_for_preview(file_path)
return image, page_info
def handle_page_turn(direction):
"""Handle page navigation"""
image, page_info, result = turn_page(direction)
return image, page_info, result
def clear_all():
"""Clear all data and reset interface"""
global pdf_cache, processing_results
pdf_cache = {
"images": [],
"current_page": 0,
"total_pages": 0,
"file_type": None,
"is_parsed": False,
"results": []
}
processing_results = {
'original_image': None,
'processed_image': None,
'layout_result': None,
'markdown_content': None,
'raw_output': None,
}
return (
None, # file_input
None, # image_preview
"No file loaded", # page_info
None, # processed_image
"Click 'Process Document' to see extracted content...", # markdown_output
"", # raw_output
None, # json_output
'✅ Interface cleared
' # model_status
)
# Wire up event handlers
demo.load(load_model_on_startup, outputs=[model_status])
file_input.change(
handle_file_upload,
inputs=[file_input],
outputs=[image_preview, page_info]
)
prev_page_btn.click(
lambda: handle_page_turn("prev"),
outputs=[image_preview, page_info, markdown_output]
)
next_page_btn.click(
lambda: handle_page_turn("next"),
outputs=[image_preview, page_info, markdown_output]
)
prompt_mode.change(
update_prompt_display,
inputs=[prompt_mode],
outputs=[prompt_display]
)
process_btn.click(
process_document,
inputs=[file_input, prompt_mode, max_new_tokens, min_pixels, max_pixels],
outputs=[processed_image, markdown_output, raw_output, json_output, model_status]
)
clear_btn.click(
clear_all,
outputs=[
file_input, image_preview, page_info, processed_image,
markdown_output, raw_output, json_output, model_status
]
)
return demo
if __name__ == "__main__":
# Create and launch the interface
demo = create_gradio_interface()
demo.queue(max_size=10).launch(
server_name="0.0.0.0",
server_port=7860,
share=False,
debug=True,
show_error=True
)