Update app.py

app.py

@@ -129,19 +129,13 @@ def generate_annotated_image(dataset_name, split, domain, category):
       example = dict()
 
   if len(msg_error) > 0:
-    # save PDF
-    rgba = Image.open(images_wo_content)
-    rgb = Image.new('RGB', rgba.size, (255, 255, 255))  # white background
-    rgb.paste(rgba, mask=rgba.split()[3])               # paste using alpha channel as mask
-    rgb.save("wo_content.pdf", 'PDF', resolution=100.0)
     # save image files
-    Image.open(images_wo_content).save("img_paragraphs.png")
-    Image.open(images_wo_content).save("img_lines.png")
+    Image.open(images_wo_content).save("wo_content.png")
     # save csv files
-    df_paragraphs_wo_content.to_csv("paragraphs.csv", encoding="utf-8", index=False)
-    df_lines_wo_content.to_csv("lines.csv", encoding="utf-8", index=False)
+    df_paragraphs_wo_content.to_csv("paragraphs_wo_content.csv", encoding="utf-8", index=False)
+    df_lines_wo_content.to_csv("lines_wo_content.csv", encoding="utf-8", index=False)
 
-    return msg_error, "wo_content.pdf", images_wo_content, images_wo_content, "img_paragraphs.png", "img_lines.png", df_paragraphs_wo_content, df_lines_wo_content, gr.File.update(value="paragraphs.csv", visible=False), gr.File.update(value="lines.csv", visible=False)
+    return msg_error, "wo_content.png", images_wo_content, images_wo_content, "wo_content.png", "wo_content.png", df_paragraphs_wo_content, df_lines_wo_content, gr.File.update(value="paragraphs_wo_content.csv", visible=False), gr.File.update(value="lines_wo_content.csv", visible=False)
   else:
     # get random image & PDF data
     index = random.randint(0, len(example))
@@ -155,8 +149,9 @@ def generate_annotated_image(dataset_name, split, domain, category):
     # resize image to original
     image = image.resize((original_width, original_height))
 
-    # get pdf of image
-    image.save(original_filename)
+    # get image of PDF without bounding boxes
+    img_file = original_filename.replace(".pdf", ".png")
+    image.save(img_file)
 
     # get corresponding annotations
     texts = example[index]["texts"]
@@ -233,8 +228,9 @@ def generate_annotated_image(dataset_name, split, domain, category):
         if i == 0: 
           imgs["paragraphs"] = img
 
-          # save 
-          img.save("img_paragraphs.png")
+          # save
+          img_paragraphs = "img_paragraphs_" + original_filename.replace(".pdf", ".png")
+          img.save(img_paragraphs)
         
           df_paragraphs["paragraphs"] = list(range(len(sorted_original_bboxes_block_list)))
           df_paragraphs["categories"] = [id2label[label_idx] for label_idx in sorted_category_block_list]
@@ -242,13 +238,15 @@ def generate_annotated_image(dataset_name, split, domain, category):
           df_paragraphs["bounding boxes"] = [str(bbox) for bbox in sorted_original_bboxes_block_list]
 
           # save 
-          df_paragraphs.to_csv("paragraphs.csv", encoding="utf-8", index=False)
+          csv_paragraphs = "csv_paragraphs_" + original_filename.replace(".pdf", ".csv")
+          df_paragraphs.to_csv(csv_paragraphs, encoding="utf-8", index=False)
 
         else: 
           imgs["lines"] = img
 
           # save 
-          img.save("img_lines.png")
+          img_lines = "img_lines_" + original_filename.replace(".pdf", ".png")
+          img.save(img_lines)
 
           df_lines["lines"] = list(range(len(sorted_original_bboxes_line_list)))
           df_lines["categories"] = [id2label[label_idx] for label_idx in sorted_category_line_list]
@@ -256,11 +254,12 @@ def generate_annotated_image(dataset_name, split, domain, category):
           df_lines["bounding boxes"] = [str(bbox) for bbox in sorted_original_bboxes_line_list]
           
           # save 
-          df_lines.to_csv("lines.csv", encoding="utf-8", index=False)
+          csv_lines = "csv_lines_" + original_filename.replace(".pdf", ".csv")
+          df_lines.to_csv(csv_lines, encoding="utf-8", index=False)
 
     msg = f'The page {page_no} of the PDF "{original_filename}" (domain: "{domain}") matches your settings.'
 
-    return msg, original_filename, imgs["paragraphs"], imgs["lines"], "img_paragraphs.png", "img_lines.png", df_paragraphs, df_lines, gr.File.update(value="paragraphs.csv", visible=True), gr.File.update(value="lines.csv", visible=True)
+    return msg, img_file, imgs["paragraphs"], imgs["lines"], img_paragraphs, img_lines, df_paragraphs, df_lines, gr.File.update(value=csv_paragraphs, visible=True), gr.File.update(value=csv_lines, visible=True)
 
 # gradio APP
 with gr.Blocks(title="DocLayNet image viewer", css=".gradio-container") as demo:
@@ -270,8 +269,7 @@ with gr.Blocks(title="DocLayNet image viewer", css=".gradio-container") as demo:
     <div><p>It uses the datasets <a style="text-decoration: none; border-bottom: #64b5f6 0.125em solid; color: #64b5f6" href="https://huggingface.co/datasets/pierreguillou/DocLayNet-small" target="_blank">DocLayNet small</a> and <a style="text-decoration: none; border-bottom: #64b5f6 0.125em solid; color: #64b5f6" style="text-decoration: none; border-bottom: #64b5f6 0.125em solid; color: #64b5f6" href="https://huggingface.co/datasets/pierreguillou/DocLayNet-base" target="_blank">DocLayNet base</a> (you can also run this APP in Google Colab by running this <a style="text-decoration: none; border-bottom: #64b5f6 0.125em solid; color: #64b5f6" href="https://github.com/piegu/language-models/blob/master/DocLayNet_image_viewer_APP.ipynb" target="_blank">notebook</a>).</p></div>
     <div><p>Make your settings and the output will show 2 images of a randomly selected PDF with labeled bounding boxes, one of paragraphs and the other of lines, and their corresponding tables of texts with their labels.</p></div>
     <div><p>For example, if you select the domain "laws_and_regulations" and the category "Caption", you will get a random PDF that corresponds to these settings (ie, it will have at least one bounding box labeled with "Caption" in the PDF).</p></div>
-    <div style="margin-top: 20px"><p>More information about the DocLayNet datasets and this APP in the 2 following blog posts:</p></div>
-    <div><ul><li>- <a style="text-decoration: none; border-bottom: #64b5f6 0.125em solid; color: #64b5f6" href="https://medium.com/@pierre_guillou/document-ai-processing-of-doclaynet-dataset-to-be-used-by-layout-models-of-the-hugging-face-hub-308d8bd81cdb" target="_blank">(01/27/2023) Document AI | Processing of DocLayNet dataset to be used by layout models of the Hugging Face hub (finetuning, inference)</a></li><li>- <a href="" target="_blank">(01/31/2023) Document AI | DocLayNet image viewer APP</a></li></ul></div>
+    <div style="margin-top: 20px"><p>More information about the DocLayNet datasets and this APP in the following blog post: <a style="text-decoration: none; border-bottom: #64b5f6 0.125em solid; color: #64b5f6" href="https://medium.com/@pierre_guillou/document-ai-processing-of-doclaynet-dataset-to-be-used-by-layout-models-of-the-hugging-face-hub-308d8bd81cdb" target="_blank">(01/27/2023) Document AI | Processing of DocLayNet dataset to be used by layout models of the Hugging Face hub (finetuning, inference)</a></div>
     """)
     with gr.Row():
         with gr.Column():
@@ -287,7 +285,7 @@ with gr.Blocks(title="DocLayNet image viewer", css=".gradio-container") as demo:
       with gr.Column():
         output_msg = gr.Textbox(label="Output message")
       with gr.Column():
-        pdf_file = gr.File(visible=True, label="PDF file (original)")
+        img_file = gr.File(visible=True, label="Image file of the PDF")
     with gr.Row():
         with gr.Column():
           img_paragraphs_file = gr.File(visible=True, label="Image file (labeled paragraphs)")
@@ -322,15 +320,15 @@ with gr.Blocks(title="DocLayNet image viewer", css=".gradio-container") as demo:
               type="pandas",
               wrap=True
             )
-    btn.click(generate_annotated_image, inputs=[dataset_name_gr, split_gr, domain_gr, category_gr], outputs=[output_msg, pdf_file, img_paragraphs, img_lines, img_paragraphs_file, img_lines_file, df_paragraphs, df_lines, csv_paragraphs, csv_lines])
+    btn.click(generate_annotated_image, inputs=[dataset_name_gr, split_gr, domain_gr, category_gr], outputs=[output_msg, img_file, img_paragraphs, img_lines, img_paragraphs_file, img_lines_file, df_paragraphs, df_lines, csv_paragraphs, csv_lines])
 
     gr.Markdown("## Example")
     gr.Examples(
         [["small", "all", "all", "all"]],
         [dataset_name_gr, split_gr, domain_gr, category_gr],
-        [output_msg, pdf_file, img_paragraphs, img_lines, img_paragraphs_file, img_lines_file, df_paragraphs, df_lines, csv_paragraphs, csv_lines],
+        [output_msg, img_file, img_paragraphs, img_lines, img_paragraphs_file, img_lines_file, df_paragraphs, df_lines, csv_paragraphs, csv_lines],
         fn=generate_annotated_image,
         cache_examples=True,
     )
 
-demo.launch()      
+demo.launch()