nougat / app.py
tomriddle's picture
Update app.py
b064a4e
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)