Spaces:
Running
on
T4
Running
on
T4
| import gradio as gr | |
| import subprocess | |
| import uuid | |
| import os | |
| import requests | |
| import re | |
| def get_pdf(pdf_link): | |
| # Generate a unique filename | |
| unique_filename = f"input/downloaded_paper_{uuid.uuid4().hex}.pdf" | |
| # Send a GET request to the PDF link | |
| response = requests.get(pdf_link) | |
| if response.status_code == 200: | |
| # Save the PDF content to a local file | |
| with open(unique_filename, 'wb') as pdf_file: | |
| pdf_file.write(response.content) | |
| print("PDF downloaded successfully.") | |
| else: | |
| print("Failed to download the PDF.") | |
| return unique_filename #.split('/')[-1][:-4] | |
| def nougat_ocr(file_name): | |
| #unique_filename = f"/content/output/downloaded_paper_{uuid.uuid4().hex}.pdf" | |
| # Command to run | |
| cli_command = [ | |
| 'nougat', | |
| #'--out', unique_filename, | |
| '--out', 'output', | |
| 'pdf', f'{file_name}', | |
| '--checkpoint', 'nougat', | |
| '--markdown' | |
| ] | |
| # Run the command and capture its output | |
| #completed_process = | |
| subprocess.run(cli_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) | |
| return #unique_filename | |
| def predict(pdf_file, pdf_link): | |
| if pdf_file is None: | |
| if pdf_link == '': | |
| print("No file is uploaded and No link is provided") | |
| return "No data provided. Upload a pdf file or provide a pdf link and try again!" | |
| else: | |
| print(f'pdf_link is - {pdf_link}') | |
| file_name = get_pdf(pdf_link) | |
| print(f'file_name is - {file_name}') | |
| else: | |
| file_name = pdf_file.name | |
| print(file_name) | |
| pdf_name = pdf_file.name.split('/')[-1].split('.')[0] | |
| print(pdf_name) | |
| # Call nougat | |
| nougat_ocr(file_name) | |
| #print("BACKKKK") | |
| # Open the file for reading | |
| file_name = file_name.split('/')[-1][:-4] | |
| with open(f'output/{file_name}.mmd', 'r') as file: | |
| content = file.read() | |
| # switch math delimiters | |
| content = content.replace(r'\(', '$').replace(r'\)', '$').replace(r'\[', '$$').replace(r'\]', '$$') | |
| return content | |
| def nougat_ocr1(file_name): | |
| print('******* inside nougat_ocr *******') | |
| # CLI Command to run | |
| cli_command = [ | |
| 'nougat', | |
| '--out', 'output', | |
| 'pdf', f'{file_name}', | |
| '--checkpoint', 'nougat', | |
| '--markdown' | |
| ] | |
| # Run the command and get .mmd file in an output folder | |
| subprocess.run(cli_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) | |
| return | |
| def predict1(pdf_file): | |
| print('******* inside predict *******') | |
| print(f"temporary file - {pdf_file.name}") | |
| pdf_name = pdf_file.name.split('/')[-1].split('.')[0] | |
| print(f"pdf file name - {pdf_name}") | |
| #! Get prediction for a PDF using nougat | |
| nougat_ocr(pdf_file.name) | |
| print("BAACCKKK") | |
| # Open the multimarkdown (.mmd) file for reading | |
| with open(f'output/{pdf_name}.mmd', 'r') as file: | |
| content = file.read() | |
| return content | |
| def process_example(pdf_file,pdf_link): | |
| ocr_content = predict(pdf_file,pdf_link) | |
| return gr.update(value=ocr_content) | |
| css = """ | |
| #mkd { | |
| height: 500px; | |
| overflow: auto; | |
| border: 1px solid #ccc; | |
| } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| gr.HTML("<h1><center>Nougat: Neural Optical Understanding for Academic Documents<center><h1>") | |
| gr.HTML("<h3><center>Lukas Blecher et al. <a href='https://arxiv.org/pdf/2308.13418.pdf' target='_blank'>Paper</a>, <a href='https://facebookresearch.github.io/nougat/'>Project</a><center></h3>") | |
| with gr.Row(): | |
| mkd = gr.Markdown('<h4><center>Upload a PDF</center></h4>',scale=1) | |
| mkd = gr.Markdown('<h4><center><i>OR</i></center></h4>',scale=1) | |
| mkd = gr.Markdown('<h4><center>Provide a PDF link</center></h4>',scale=1) | |
| with gr.Row(equal_height=True): | |
| pdf_file = gr.File(label='PDF📃', file_count='single', scale=1) | |
| pdf_link = gr.Textbox(placeholder='Enter an Arxiv link here', label='PDF link🔗🌐', scale=1) | |
| with gr.Row(): | |
| btn = gr.Button('Run NOUGAT🍫') | |
| clr = gr.Button('Clear🚿') | |
| output_headline = gr.Markdown("<h3>PDF converted to markup language through Nougat-OCR👇:</h3>") | |
| parsed_output = gr.Markdown(elem_id='mkd', value='📃🔤OCR Output') | |
| btn.click(predict, [pdf_file, pdf_link], parsed_output ) | |
| clr.click(lambda : (gr.update(value=None), | |
| gr.update(value=None), | |
| gr.update(value=None)), | |
| [], | |
| [pdf_file, pdf_link, parsed_output] | |
| ) | |
| gr.Examples( | |
| [["input/nougat.pdf", ""], [None, "https://arxiv.org/pdf/2308.08316.pdf"]], | |
| inputs = [pdf_file, pdf_link], | |
| outputs = parsed_output, | |
| fn=process_example, | |
| cache_examples=True, | |
| label='Click on any Examples below to get Nougat OCR results quickly:' | |
| ) | |
| demo.queue() | |
| demo.launch(debug=True) | |