| import os |
| from doctr.io import DocumentFile |
| from doctr.models import ocr_predictor, from_hub |
| import gradio as gr |
|
|
| os.environ['USE_TORCH'] = '1' |
| reco_model_zgh = from_hub('ayymen/crnn_mobilenet_v3_large_zgh') |
| predictor_zgh = ocr_predictor(reco_arch=reco_model_zgh, pretrained=True) |
|
|
| reco_model = from_hub('ayymen/crnn_mobilenet_v3_large_tifinagh') |
| predictor = ocr_predictor(reco_arch=reco_model, pretrained=True) |
|
|
| title = "Tifinagh OCR" |
| description = """Upload an image to get the OCR results! |
| """ |
|
|
| def ocr(img, script): |
| img.save("out.jpg") |
| doc = DocumentFile.from_images("out.jpg") |
| output = predictor_zgh(doc) if script == "Tifinagh-IRCAM" else predictor(doc) |
| res = "" |
| for obj in output.pages: |
| for obj1 in obj.blocks: |
| for obj2 in obj1.lines: |
| for obj3 in obj2.words: |
| res = res + " " + obj3.value |
| res = res + "\n" |
| res = res + "\n" |
| _output_name = "RESULT_OCR.txt" |
| open(_output_name, 'w', encoding="utf-8").close() |
| with open(_output_name, "w", encoding="utf-8", errors="ignore") as f: |
| f.write(res) |
| print("Writing into file") |
| return res, _output_name |
|
|
| demo = gr.Interface(fn=ocr, |
| inputs=[ |
| gr.Image(type="pil"), |
| gr.Dropdown(choices=['Tifinagh-IRCAM', 'Tifinagh'], label="Script", value="Tifinagh-IRCAM") |
| ], |
| outputs=[ |
| gr.Textbox(lines=20, label="Full Text"), |
| gr.File(label="Download OCR Results") |
| ], |
| title=title, |
| description=description, |
| examples=[ |
| ["Examples/3.jpg", "Tifinagh-IRCAM"], |
| ["Examples/2.jpg", "Tifinagh-IRCAM"], |
| ["Examples/1.jpg", "Tifinagh-IRCAM"] |
| ] |
| ) |
|
|
| demo.launch(debug=True) |
|
|