config.py

"""
Configuration file for dots.ocr Space
"""

# Model configuration
MODEL_PATH = "rednote-hilab/DotsOCR"

# Generation parameters
MAX_NEW_TOKENS = 8192
TEMPERATURE = 0.1
TOP_P = 0.9

# GPU configuration
GPU_DURATION = 120  # seconds for @spaces.GPU decorator

# Image constraints
MAX_IMAGE_PIXELS = 11289600  # Maximum recommended pixels for best performance
RECOMMENDED_DPI = 200  # For PDF conversion

# Supported image formats
SUPPORTED_FORMATS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
SUPPORTED_DOC_FORMATS = ['.pdf']

# Prompt templates
PROMPT_TEMPLATES = {
    "full_layout_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.""",

    "ocr_only": """Please extract all text from the image in reading order. Format the output as plain text, preserving the original structure as much as possible.""",
    
    "layout_only": """Please detect all layout elements in the image and output their bounding boxes and categories. Format: [{"bbox": [x1, y1, x2, y2], "category": "category_name"}]""",
}

# UI Configuration
UI_TITLE = "dots.ocr - Multilingual Document OCR"
UI_DESCRIPTION = """
Upload a document image and get OCR results with layout detection.
This space uses the [dots.ocr](https://github.com/rednote-hilab/dots.ocr) model.
"""

LAYOUT_CATEGORIES = [
    'Caption', 'Footnote', 'Formula', 'List-item', 
    'Page-footer', 'Page-header', 'Picture', 
    'Section-header', 'Table', 'Text', 'Title'
]