| import gradio as gr |
| import torch |
| import spaces |
| import os |
| import tempfile |
| from PIL import Image, ImageOps |
| from typing import Iterable |
|
|
| from transformers import AutoProcessor, AutoModelForImageTextToText |
| from gradio.themes import Soft |
| from gradio.themes.utils import colors, fonts, sizes |
|
|
| colors.hot_pink = colors.Color( |
| name="hot_pink", |
| c50="#FFF0F5", |
| c100="#FFE4EC", |
| c200="#FFC0D9", |
| c300="#FF99C4", |
| c400="#FF7EB8", |
| c500="#FF69B4", |
| c600="#E55AA0", |
| c700="#CC4C8C", |
| c800="#B33D78", |
| c900="#992F64", |
| c950="#802050", |
| ) |
|
|
|
|
| class HotPinkTheme(Soft): |
| def __init__( |
| self, |
| *, |
| primary_hue: colors.Color | str = colors.gray, |
| secondary_hue: colors.Color | str = colors.hot_pink, |
| neutral_hue: colors.Color | str = colors.slate, |
| text_size: sizes.Size | str = sizes.text_lg, |
| font: fonts.Font | str | Iterable[fonts.Font | str] = ( |
| fonts.GoogleFont("Outfit"), |
| "Arial", |
| "sans-serif", |
| ), |
| font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( |
| fonts.GoogleFont("IBM Plex Mono"), |
| "ui-monospace", |
| "monospace", |
| ), |
| ): |
| super().__init__( |
| primary_hue=primary_hue, |
| secondary_hue=secondary_hue, |
| neutral_hue=neutral_hue, |
| text_size=text_size, |
| font=font, |
| font_mono=font_mono, |
| ) |
| super().set( |
| background_fill_primary="*primary_50", |
| background_fill_primary_dark="*primary_900", |
| body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", |
| body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", |
| button_primary_text_color="white", |
| button_primary_text_color_hover="white", |
| button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", |
| button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", |
| button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", |
| button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", |
| button_secondary_text_color="black", |
| button_secondary_text_color_hover="white", |
| button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)", |
| button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)", |
| button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)", |
| button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)", |
| slider_color="*secondary_500", |
| slider_color_dark="*secondary_600", |
| block_title_text_weight="600", |
| block_border_width="3px", |
| block_shadow="*shadow_drop_lg", |
| button_primary_shadow="*shadow_drop_lg", |
| button_large_padding="11px", |
| color_accent_soft="*primary_100", |
| block_label_background_fill="*primary_200", |
| ) |
|
|
|
|
| hot_pink_theme = HotPinkTheme() |
|
|
| css = """ |
| @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500;600&display=swap'); |
| |
| /* Grid background */ |
| body, .gradio-container { |
| background-color: #FFF0F5 !important; |
| background-image: |
| linear-gradient(#FFC0D9 1px, transparent 1px), |
| linear-gradient(90deg, #FFC0D9 1px, transparent 1px) !important; |
| background-size: 40px 40px !important; |
| font-family: 'Outfit', sans-serif !important; |
| } |
| |
| .dark body, .dark .gradio-container { |
| background-color: #1a1a1a !important; |
| background-image: |
| linear-gradient(rgba(255, 105, 180, 0.1) 1px, transparent 1px), |
| linear-gradient(90deg, rgba(255, 105, 180, 0.1) 1px, transparent 1px) !important; |
| background-size: 40px 40px !important; |
| } |
| |
| /* Sidebar width */ |
| .gradio-sidebar { |
| min-width: 420px !important; |
| max-width: 480px !important; |
| } |
| |
| /* Titles */ |
| #main-title h1 { |
| font-size: 2.5em !important; |
| font-weight: 700 !important; |
| background: linear-gradient(135deg, #FF69B4 0%, #FF99C4 50%, #E55AA0 100%); |
| background-size: 200% 200%; |
| -webkit-background-clip: text; |
| -webkit-text-fill-color: transparent; |
| background-clip: text; |
| animation: gradient-shift 4s ease infinite; |
| letter-spacing: -0.02em; |
| } |
| |
| #output-title h2 { |
| font-size: 2.2em !important; |
| } |
| |
| @keyframes gradient-shift { |
| 0%, 100% { background-position: 0% 50%; } |
| 50% { background-position: 100% 50%; } |
| } |
| |
| /* Card styling */ |
| .gradio-group { |
| background: rgba(255, 255, 255, 0.9) !important; |
| border: 2px solid #FFC0D9 !important; |
| border-radius: 12px !important; |
| box-shadow: 0 4px 24px rgba(255, 105, 180, 0.08) !important; |
| backdrop-filter: blur(10px); |
| transition: all 0.3s ease; |
| } |
| |
| .gradio-group:hover { |
| box-shadow: 0 8px 32px rgba(255, 105, 180, 0.12) !important; |
| border-color: #FF99C4 !important; |
| } |
| |
| .dark .gradio-group { |
| background: rgba(30, 30, 30, 0.9) !important; |
| border-color: rgba(255, 105, 180, 0.3) !important; |
| } |
| |
| /* Image upload */ |
| .gradio-image { |
| border-radius: 10px !important; |
| overflow: hidden; |
| border: 2px dashed #FF99C4 !important; |
| transition: all 0.3s ease; |
| } |
| |
| .gradio-image:hover { |
| border-color: #FF69B4 !important; |
| background: rgba(255, 105, 180, 0.02) !important; |
| } |
| |
| /* Radio buttons */ |
| .gradio-radio label { |
| border-radius: 6px !important; |
| transition: all 0.2s ease !important; |
| border: 1px solid transparent !important; |
| } |
| |
| .gradio-radio label:hover { |
| background: rgba(255, 105, 180, 0.05) !important; |
| } |
| |
| .gradio-radio label.selected { |
| background: rgba(255, 105, 180, 0.1) !important; |
| border-color: #FF69B4 !important; |
| } |
| |
| /* Primary button */ |
| .primary { |
| border-radius: 8px !important; |
| font-weight: 600 !important; |
| letter-spacing: 0.05em !important; |
| transition: all 0.3s ease !important; |
| } |
| |
| .primary:hover { |
| transform: translateY(-2px) !important; |
| } |
| |
| /* Output textbox */ |
| .gradio-textbox textarea { |
| font-family: 'IBM Plex Mono', monospace !important; |
| font-size: 0.95rem !important; |
| line-height: 1.7 !important; |
| background: rgba(255, 255, 255, 0.95) !important; |
| border: 1px solid #FFC0D9 !important; |
| border-radius: 8px !important; |
| } |
| |
| .dark .gradio-textbox textarea { |
| background: rgba(30, 30, 30, 0.95) !important; |
| border-color: rgba(255, 105, 180, 0.2) !important; |
| } |
| |
| /* Markdown output */ |
| .gradio-markdown { |
| font-family: 'Outfit', sans-serif !important; |
| line-height: 1.7 !important; |
| } |
| |
| .gradio-markdown code { |
| font-family: 'IBM Plex Mono', monospace !important; |
| background: rgba(255, 105, 180, 0.08) !important; |
| padding: 2px 6px !important; |
| border-radius: 4px !important; |
| color: #CC4C8C !important; |
| } |
| |
| .gradio-markdown pre { |
| background: rgba(255, 105, 180, 0.05) !important; |
| border: 1px solid #FFC0D9 !important; |
| border-radius: 8px !important; |
| padding: 1rem !important; |
| } |
| |
| /* Examples */ |
| .gradio-examples .gallery-item { |
| border: 2px solid #FFC0D9 !important; |
| border-radius: 8px !important; |
| transition: all 0.2s ease !important; |
| } |
| |
| .gradio-examples .gallery-item:hover { |
| border-color: #FF69B4 !important; |
| transform: translateY(-2px) !important; |
| box-shadow: 0 4px 12px rgba(255, 105, 180, 0.15) !important; |
| } |
| |
| /* Scrollbar */ |
| ::-webkit-scrollbar { width: 8px; height: 8px; } |
| ::-webkit-scrollbar-track { background: rgba(255,105,180,0.05); border-radius: 4px; } |
| ::-webkit-scrollbar-thumb { background: linear-gradient(135deg, #FF69B4, #FF99C4); border-radius: 4px; } |
| ::-webkit-scrollbar-thumb:hover { background: linear-gradient(135deg, #E55AA0, #FF69B4); } |
| |
| /* Accordion */ |
| .gradio-accordion { |
| border-radius: 10px !important; |
| border: 1px solid #FFC0D9 !important; |
| } |
| |
| .gradio-accordion > .label-wrap { |
| background: rgba(255, 105, 180, 0.03) !important; |
| border-radius: 10px !important; |
| } |
| |
| /* Animations */ |
| @keyframes fadeIn { |
| from { opacity: 0; transform: translateY(10px); } |
| to { opacity: 1; transform: translateY(0); } |
| } |
| |
| .gradio-row { animation: fadeIn 0.4s ease-out; } |
| |
| label { font-weight: 600 !important; color: #333 !important; } |
| .dark label { color: #eee !important; } |
| |
| footer { display: none !important; } |
| |
| |
| /* Wider sidebar */ |
| .sidebar { |
| min-width: 420px !important; |
| max-width: 480px !important; |
| } |
| """ |
|
|
| MODEL_PATH = "zai-org/GLM-OCR" |
|
|
| processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True) |
| model = AutoModelForImageTextToText.from_pretrained( |
| pretrained_model_name_or_path=MODEL_PATH, |
| torch_dtype=torch.bfloat16, |
| device_map="auto", |
| trust_remote_code=True, |
| ) |
|
|
| TASK_PROMPTS = { |
| "Text": "Text Recognition:", |
| "Formula": "Formula Recognition:", |
| "Table": "Table Recognition:", |
| } |
|
|
| @spaces.GPU |
| def process_image(image, task): |
| """Run OCR on the uploaded image with the selected recognition type.""" |
| if image is None: |
| return "Please upload an image first.", "Please upload an image first." |
|
|
| if image.mode in ("RGBA", "LA", "P"): |
| image = image.convert("RGB") |
| image = ImageOps.exif_transpose(image) |
|
|
| tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png") |
| image.save(tmp.name, "PNG") |
| tmp.close() |
|
|
| prompt = TASK_PROMPTS.get(task, "Text Recognition:") |
|
|
| messages = [ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image", "url": tmp.name}, |
| {"type": "text", "text": prompt}, |
| ], |
| } |
| ] |
|
|
| inputs = processor.apply_chat_template( |
| messages, |
| tokenize=True, |
| add_generation_prompt=True, |
| return_dict=True, |
| return_tensors="pt", |
| ).to(model.device) |
|
|
| inputs.pop("token_type_ids", None) |
|
|
| generated_ids = model.generate(**inputs, max_new_tokens=8192) |
| output_text = processor.decode( |
| generated_ids[0][inputs["input_ids"].shape[1] :], |
| skip_special_tokens=True, |
| ) |
|
|
| os.unlink(tmp.name) |
|
|
| result = output_text.strip() |
| return result, result |
|
|
| with gr.Blocks(fill_height=True) as demo: |
|
|
| with gr.Sidebar(width=450): |
|
|
| gr.Markdown("# **GLM-OCR**", elem_id="main-title") |
|
|
| image_input = gr.Image( |
| type="pil", |
| label="Upload Image", |
| sources=["upload", "clipboard"], |
| height=300, |
| ) |
|
|
| task = gr.Radio( |
| choices=list(TASK_PROMPTS.keys()), |
| value="Text", |
| label="Recognition Type", |
| ) |
|
|
| btn = gr.Button("Perform OCR", variant="primary") |
|
|
| gr.Examples( |
| examples=[ |
| "examples/1.jpg", |
| "examples/4.jpg", |
| "examples/5.webp", |
| "examples/2.jpg", |
| "examples/3.jpg", |
| ], |
| inputs=image_input, |
| label="Examples", |
| ) |
|
|
| gr.Markdown("## Output", elem_id="output-title") |
|
|
| output_text = gr.Textbox( |
| label="Raw Output Stream", |
| interactive=True, |
| lines=22, |
| ) |
|
|
| with gr.Accordion("(Result.md)", open=False): |
| output_md = gr.Markdown(label="Rendered Markdown") |
|
|
| btn.click( |
| fn=process_image, |
| inputs=[image_input, task], |
| outputs=[output_text, output_md], |
| ) |
|
|
| image_input.change( |
| fn=lambda: ("", ""), |
| inputs=None, |
| outputs=[output_text, output_md], |
| ) |
|
|
| if __name__ == "__main__": |
| demo.queue(max_size=50).launch( |
| css=css, |
| theme=hot_pink_theme, |
| mcp_server=True, |
| ssr_mode=False, |
| show_error=True, |
| ) |
