Update app.py

app.py

@@ -1,13 +1,32 @@
-import gradio as gr
-import easyocr
+# import gradio as gr
+# import easyocr
+
+# reader = easyocr.Reader(["en"])
+
+# def ocr_image(image):
+#     results = reader.readtext(image)
+#     return "\n".join([text for _, text, _ in results])
 
-reader = easyocr.Reader(["en"])
+# demo = gr.Interface(fn=ocr_image, inputs="image", outputs="text")
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+from PIL import Image
+from huggingface_hub import snapshot_download
+from transformers import AutoProcessor, AutoModelForCausalLM
+import gradio as gr
+import torch
 
-def ocr_image(image):
-    results = reader.readtext(image)
-    return "\n".join([text for _, text, _ in results])
+MODEL_ID = "rednote-hilab/dots.ocr"
+local = snapshot_download(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(local, trust_remote_code=True, device_map="auto", torch_dtype=torch.bfloat16)
+processor = AutoProcessor.from_pretrained(local, trust_remote_code=True)
 
-demo = gr.Interface(fn=ocr_image, inputs="image", outputs="text")
+def parse_document(image):
+    inputs = processor(images=[image], return_tensors="pt").to(model.device)
+    output = model.generate(**inputs, do_sample=False, max_new_tokens=1024)
+    return processor.batch_decode(output, skip_special_tokens=True)[0]
 
-if __name__ == "__main__":
-    demo.launch()
+demo = gr.Interface(parse_document, inputs=gr.Image(type="pil"), outputs="text")
+if __name__ == "__main__": demo.launch()