Update app.py

app.py

@@ -1,83 +1,42 @@
-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
-
-
-def nougat_ocr(file_name):
-    # Command to run
-    cli_command = [
-        'nougat',
-        '--out', 'output',
-        'pdf', f'{file_name}',
-        '--checkpoint', 'nougat',
-        '--markdown'
-    ]
-
-    # Run the command and capture its output
-    subprocess.run(cli_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-
-
-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)
-
-    # Open the file for reading
-    file_name = file_name.split('/')[-1][:-4]
-    mmd_file_path = f'output/{file_name}.mmd'
-    with open(mmd_file_path, 'r') as file:
-        content = file.read()
-    # switch math delimiters
-    content = content.replace(r'\(', '$').replace(r'\)', '$').replace(r'\[', '$$').replace(r'\]', '$$')
-    
-    return content, mmd_file_path
-
-
-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><
+  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')
+  mmd_file_download = gr.File(label='Download .mmd file', interactive=False)
+
+  def handle_predict(pdf_file, pdf_link):
+      content, mmd_file_path = predict(pdf_file, pdf_link)
+      return gr.update(value=content), mmd_file_path
+
+  btn.click(handle_predict, [pdf_file, pdf_link], [parsed_output, mmd_file_download])
+  clr.click(lambda: (gr.update(value=None),
+                     gr.update(value=None),
+                     gr.update(value=None)),
+            [],
+            [pdf_file, pdf_link, parsed_output, mmd_file_download])
+
+  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)