Update app.py

app.py

@@ -1,60 +1,64 @@
 import gradio as gr
 import spaces
 from transformers import AutoModel, AutoTokenizer
+from PIL import Image
+import numpy as np
 
 tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
 model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)
 model = model.eval().cuda()
 
 @spaces.GPU
-def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=True):
-    if task == "Plain Text OCR":
-        res = model.chat(tokenizer, image, ocr_type='ocr')
-    elif task == "Format Text OCR":
-        res = model.chat(tokenizer, image, ocr_type='format')
-    elif task == "Fine-grained OCR (Box)":
-        res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box)
-    elif task == "Fine-grained OCR (Color)":
-        res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color)
-    elif task == "Multi-crop OCR":
-        res = model.chat_crop(tokenizer, image_file=image)
-    elif task == "Render Formatted OCR":
-        res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html')
+def run_GOT(image_array, got_mode, ocr_box="", ocr_color=""):
+    image = Image.fromarray(np.uint8(image_array))
+    if got_mode == "plain texts OCR":
+        res = model.chat(tokenizer, image, ocr_type='ocr', gradio_input=True)
+    elif got_mode == "format texts OCR":
+        res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html', gradio_input=True)
+    elif got_mode == "plain multi-crop OCR":
+        res = model.chat_crop(tokenizer, image, ocr_type='ocr', gradio_input=True)
+    elif got_mode == "format multi-crop OCR":
+        res = model.chat_crop(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html', gradio_input=True)
+
+    elif got_mode == "plain fine-grained OCR":
+        res = model.chat(tokenizer, image, ocr_type='ocr', ocr_box=ocr_box, ocr_color=ocr_color, gradio_input=True)
+    elif got_mode == "format fine-grained OCR":
+        res = model.chat(tokenizer, image, ocr_type='format', ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file='./demo.html', gradio_input=True)
+
+    if "format" in got_mode:
         with open('./demo.html', 'r') as f:
-            html_content = f.read()
-        return res, html_content
-    
+            demo_html = f.read()
+        return res, demo_html
     return res, None
 
-def update_inputs(task):
-    if task == "Plain Text OCR" or task == "Format Text OCR" or task == "Multi-crop OCR":
-        return [gr.update(visible=False)] * 4
-    elif task == "Fine-grained OCR (Box)":
+def task_update(task):
+    if "fine-grained" in task:
         return [
-            gr.update(visible=True, choices=["ocr", "format"]),
             gr.update(visible=True),
             gr.update(visible=False),
-            gr.update(visible=False)
+            gr.update(visible=False),
         ]
-    elif task == "Fine-grained OCR (Color)":
+    else:
         return [
-            gr.update(visible=True, choices=["ocr", "format"]),
             gr.update(visible=False),
-            gr.update(visible=True, choices=["red", "green", "blue"]),
-            gr.update(visible=False)
+            gr.update(visible=False),
+            gr.update(visible=False),
         ]
-    elif task == "Render Formatted OCR":
-        return [gr.update(visible=False)] * 3 + [gr.update(visible=True)]
-    
-def ocr_demo(image, task, ocr_type, ocr_box, ocr_color, render):
-    res, html_content = process_image(image, task, ocr_type, ocr_box, ocr_color, render)
-    if html_content:
-        return res, html_content
-    return res, None
+
+def fine_grained_update(task):
+    if task == "box":
+        return [
+            gr.update(visible=False, value = ""),
+            gr.update(visible=True),
+        ]
+    elif task == 'color':
+        return [
+            gr.update(visible=True),
+            gr.update(visible=False, value = ""),
+        ]
+
 
 with gr.Blocks() as demo:
-    gr.Markdown()
-    gr.Markdown()
     gr.Markdown("""
     # "General OCR Theory: Towards OCR-2.0 via a Unified End-to-end Model"
     
@@ -62,7 +66,7 @@ with gr.Blocks() as demo:
     
     ### Repo
     - **Hugging Face**: [ucaslcl/GOT-OCR2_0](https://huggingface.co/ucaslcl/GOT-OCR2_0)
-    - **GitHub**: [ucaslcl/GOT-OCR2_0](https://github.com/Ucas-HaoranWei/GOT-OCR2.0/)
+    - **GitHub**: [Ucas-HaoranWei/GOT-OCR2_0](https://github.com/Ucas-HaoranWei/GOT-OCR2.0/)
     - **Paper**: [AriXiv](https://arxiv.org/abs/2409.01704)
     """)
     
@@ -71,52 +75,64 @@ with gr.Blocks() as demo:
             image_input = gr.Image(type="filepath", label="upload your image")
             task_dropdown = gr.Dropdown(
                 choices=[
-                    "plain text OCR",
-                    "format text OCR",
-                    "fine-grained OCR (Box)",
-                    "fine-grained OCR (Color)",
-                    "multi-crop OCR",
-                    "render Formatted OCR"
+                    "plain texts OCR",
+                    "format texts OCR",
+                    "plain multi-crop OCR",
+                    "format multi-crop OCR",
+                    "plain fine-grained OCR",
+                    "format fine-grained OCR",
                 ],
-                label="Select Task",
-                value="Plain Text OCR"
+                label="Choose one mode of GOT",
+                value="plain texts OCR"
             )
-            ocr_type_dropdown = gr.Dropdown(
-                choices=["ocr", "format"],
-                label="OCR Type",
+            fine_grained_dropdown = gr.Dropdown(
+                choices=["box", "color"],
+                label="fine-grained type",
                 visible=False
             )
-            ocr_box_input = gr.Textbox(
-                label="OCR Box (x1,y1,x2,y2)",
-                placeholder="e.g., 100,100,200,200",
-                visible=False
-            )
-            ocr_color_dropdown = gr.Dropdown(
+            color_dropdown = gr.Dropdown(
                 choices=["red", "green", "blue"],
-                label="OCR Color",
+                label="color list",
                 visible=False
             )
-            render_checkbox = gr.Checkbox(
-                label="Render Result",
+            box_input = gr.Textbox(
+                label="input box: [x1,y1,x2,y2]",
+                placeholder="e.g., [0,0,100,100]",
                 visible=False
             )
-            submit_button = gr.Button("Process")
+            submit_button = gr.Button("Submit")
         
         with gr.Column():
-            output_text = gr.Textbox(label="OCR Result")
-            output_html = gr.HTML(label="Rendered HTML Output")
+            ocr_result = gr.Textbox(label="GOT output")
+            html_result = gr.HTML(label="rendered html")
+    
+    gr.Examples(
+        examples=[
+            ["assets/coco.jpg", "plain texts OCR", "", ""],
+            ["assets/en2.png", "plain texts OCR", "", ""],
+            ["assets/eq.jpg", "format texts OCR", "", ""],
+            ["assets/table.jpg", "format texts OCR", "", ""],
+            ["assets/aff2.png", "plain fine-grained OCR", "[409,763,756,891]", ""],
+        ],
+        inputs=[image_input, task_dropdown],
+        label="examples",
+    )
     
-
     task_dropdown.change(
-        update_inputs,
+        task_update,
         inputs=[task_dropdown],
-        outputs=[ocr_type_dropdown, ocr_box_input, ocr_color_dropdown, render_checkbox]
+        outputs=[fine_grained_dropdown, color_dropdown, box_input]
+    )
+    fine_grained_dropdown.change(
+        fine_grained_update,
+        inputs=[fine_grained_dropdown],
+        outputs=[color_dropdown, box_input]
     )
     
     submit_button.click(
-        ocr_demo,
-        inputs=[image_input, task_dropdown, ocr_type_dropdown, ocr_box_input, ocr_color_dropdown],
-        outputs=[output_text, output_html]
+        run_GOT,
+        inputs=[image_input, task_dropdown, box_input, color_dropdown],
+        outputs=[ocr_result, html_result]
     )
 
-demo.launch()
+demo.launch(share=True)