dots-ocr-endpoint-infernece.py

import os
from openai import OpenAI 
from dotenv import load_dotenv
import plotly.express as px
import plotly.graph_objects as go
import requests
from PIL import Image
from io import BytesIO
import numpy as np
import json

load_dotenv()
client = OpenAI(
	base_url = "",
	api_key = os.environ["HF_TOKEN"]
)

prompt = """\
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.\
"""

chat_completion = client.chat.completions.create(
	model = "rednote-hilab/dots.ocr",
	messages = [
		{
			"role": "user",
			"content": [
				{
					"type": "image_url",
					"image_url": {
						"url": "https://github.com/rednote-hilab/dots.ocr/blob/master/demo/demo_image1.jpg?raw=true"
					}
				},
				{
					"type": "text",
					"text": prompt,
				}
			]
		}
	],
	stream = True,
)

text = ""
for message in chat_completion:
	text = text + message.choices[0].delta.content

annotations = json.loads(text)

# Load image
url = "https://github.com/rednote-hilab/dots.ocr/blob/master/demo/demo_image1.jpg?raw=true"
response = requests.get(url)
img = Image.open(BytesIO(response.content))
img_array = np.array(img)

# Enhanced color mapping for different categories
category_colors = {
    'Title': '#FF6B6B',
    'Section-header': '#4ECDC4', 
    'Text': '#45B7D1',
    'Picture': '#96CEB4',
    'Table': '#FFEAA7',
    'Formula': '#DDA0DD',
    'Caption': '#98D8C8',
    'List-item': '#F7DC6F',
    'Footnote': '#BB8FCE',
    'Page-header': '#85C1E9',
    'Page-footer': '#F8C471'
}

# Create figure with enhanced settings
fig = px.imshow(img_array, aspect='equal')

# Enhanced layout configuration
fig.update_layout(
    title={
        'text': "Interactive OCR Layout Analysis",
        'x': 0.5,
        'xanchor': 'center',
        'font': {'size': 18, 'family': 'Arial Black'}
    },
    dragmode="pan",  # Enable panning
    hovermode="closest",
    margin=dict(l=20, r=20, t=60, b=20),
    showlegend=True,
    legend=dict(
        orientation="v",
        yanchor="top",
        y=1,
        xanchor="left",
        x=1.02,
        bgcolor="rgba(255,255,255,0.8)",
        bordercolor="rgba(0,0,0,0.2)",
        borderwidth=1
    ),
    plot_bgcolor='white',
    paper_bgcolor='white'
)

# Track categories for legend
added_categories = set()

# Add enhanced bounding boxes with category-based colors
for i, ann in enumerate(annotations):
    x1, y1, x2, y2 = ann['bbox']
    category = ann.get('category', 'Unknown')
    color = category_colors.get(category, '#FF4444')
    
    # Enhanced bounding box with category-specific styling
    line_width = 3 if category in ['Title', 'Section-header'] else 2
    opacity = 0.8 if category == 'Picture' else 1.0
    
    fig.add_shape(
        type="rect",
        x0=x1, y0=y1, x1=x2, y1=y2,
        line=dict(color=color, width=line_width),
        opacity=opacity
    )
    
    # Enhanced hover information with better formatting
    text_content = ann.get('text', 'No text available')
    if len(text_content) > 200:
        text_content = text_content[:200] + "..."
    
    # Format hover text based on category
    if category == 'Formula':
        hover_text = f"<b>🔢 {category}</b><br><i>{text_content}</i>"
    elif category == 'Picture':
        hover_text = f"<b>🖼️ {category}</b><br>Image element"
    elif category == 'Table':
        hover_text = f"<b>📊 {category}</b><br>{text_content}"
    elif category == 'Title':
        hover_text = f"<b>📋 {category}</b><br><b>{text_content}</b>"
    else:
        hover_text = f"<b>📄 {category}</b><br>{text_content}"
    
    # Add bbox dimensions to hover
    width = x2 - x1
    height = y2 - y1
    hover_text += f"<br><br><i>Box: {width:.0f}×{height:.0f}px</i>"
    
    # Create hover point with legend entry
    show_legend = category not in added_categories
    if show_legend:
        added_categories.add(category)
    
    fig.add_trace(go.Scatter(
        x=[(x1 + x2) / 2],
        y=[(y1 + y2) / 2],
        mode="markers",
        marker=dict(
            size=20,
            opacity=0,  # Invisible marker
            color=color
        ),
        text=[hover_text],
        hoverinfo="text",
        hovertemplate="%{text}<extra></extra>",
        name=category,
        showlegend=show_legend,
        legendgroup=category
    ))

# Enhanced axes configuration
fig.update_xaxes(
    showticklabels=False,
    showgrid=False,
    zeroline=False
)
fig.update_yaxes(
    showticklabels=False,
    showgrid=False,
    zeroline=False,
    scaleanchor="x",
    scaleratio=1
)

# Add custom controls info
fig.add_annotation(
    text="💡 Hover over colored boxes to see content • Pan: drag • Zoom: scroll",
    xref="paper", yref="paper",
    x=0.5, y=-0.05,
    showarrow=False,
    font=dict(size=12, color="gray"),
    xanchor="center"
)

# Display statistics
total_elements = len(annotations)
category_counts = {}
for ann in annotations:
    cat = ann.get('category', 'Unknown')
    category_counts[cat] = category_counts.get(cat, 0) + 1

print(f"📊 Layout Analysis Complete!")
print(f"Total elements detected: {total_elements}")
print("Category breakdown:")
for cat, count in sorted(category_counts.items()):
    print(f"  • {cat}: {count}")

fig.show()