fix: got_ocr

app.py

@@ -71,18 +71,25 @@ def on_image_select(evt: gr.SelectData):
 
 
 # 更新perform_ocr函数的输入参数
-def perform_ocr(selected_index, image_gallery, task, fine_grained_type, color, box):
+def perform_ocr(selected_index, image_gallery, got_mode, fine_grained_type, color, box):
     if selected_index is None or len(image_gallery) == 0:
         return "请先选择一张图片"
 
     selected_image = image_gallery[selected_index][0]
 
     # 根据选择的任务和参数调用GOT OCR
-    got_mode = task
     ocr_color = color if fine_grained_type == "color" else ""
     ocr_box = box if fine_grained_type == "box" else ""
 
-    result, _ = got_ocr(model, tokenizer, selected_image, got_mode=got_mode, fine_grained_mode=fine_grained_type, ocr_color=ocr_color, ocr_box=ocr_box)
+    result, _ = got_ocr(
+        model,
+        tokenizer,
+        selected_image,
+        got_mode=got_mode,
+        fine_grained_mode=fine_grained_type,
+        ocr_color=ocr_color,
+        ocr_box=ocr_box,
+    )
     return result
 
 
@@ -106,6 +113,10 @@ with gr.Blocks() as demo:
     pdf_input = gr.File(label="上传PDF文件")
     image_gallery = gr.Gallery(label="PDF页面预览", columns=3, height="auto")
     selected_index = gr.State(None)
+    pdf_input.upload(fn=process_pdf, inputs=pdf_input, outputs=image_gallery)
+    image_gallery.select(fn=on_image_select, inputs=[], outputs=selected_index)
+
+    ####################
     task_dropdown = gr.Dropdown(
         choices=["plain texts OCR", "format texts OCR", "plain multi-crop OCR", "format multi-crop OCR", "plain fine-grained OCR", "format fine-grained OCR"],
         label="选择GOT模式",
@@ -123,14 +134,11 @@ with gr.Blocks() as demo:
 
     # 更新ocr_button的click事件，传递所有必要的参数
     ocr_button.click(
-        fn=perform_ocr, inputs=[selected_index, image_gallery, task_dropdown, fine_grained_dropdown, color_dropdown, box_input], outputs=ocr_result
+        fn=perform_ocr,
+        inputs=[selected_index, image_gallery, task_dropdown, fine_grained_dropdown, color_dropdown, box_input],
+        outputs=ocr_result,
     )
 
-    pdf_input.upload(fn=process_pdf, inputs=pdf_input, outputs=image_gallery)
-    image_gallery.select(fn=on_image_select, inputs=[], outputs=selected_index)
-    ocr_button.click(fn=perform_ocr, inputs=[selected_index, image_gallery], outputs=ocr_result)
-
-    # 移除了选中图片的显示部分
 
 if __name__ == "__main__":
     demo.launch()