import gradio as gr
import pymupdf # PyMuPDF for handling PDF files
from PIL import Image
import os
from functions import get_image_informations
from dataSchema import *
# import shutil
def Noc_timeSheet_pdf_to_img(pdf_path, output_path, dpi: int = 300, quality: int = 95):
pdf_document = pymupdf.open(pdf_path)
# Get the first page of the PDF
page = pdf_document.load_page(0) # 0 is the first page
# Convert the page to a pixmap (image)
pix = page.get_pixmap(dpi=dpi)
# Convert the pixmap to a PIL Image and save as JPG
image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
width, height = image.size
start_y_total_table = int(height * 0.42)
end_y_first_table = int(height * 0.30)
croped1 = image.crop((0, 0, width // 2, end_y_first_table))
croped2 = image.crop((0, start_y_total_table, width // 2, height))
upper_width, upper_height = croped1.size
lower_width, lower_height = croped2.size
combined_image = Image.new('RGB', (upper_width, upper_height + lower_height))
# Paste the upper image (croped1) on top
combined_image.paste(croped1, (0, 0))
# Paste the lower image (croped2) below the upper image
combined_image.paste(croped2, (0, upper_height))
# Save the combined image
combined_image.save(output_path, "JPEG", quality=quality)
def Clauses_in_invoice(pdf_path: str) -> bool:
"""
Extract text from the last page of a PDF.
"""
pdf_document = pymupdf.open(pdf_path)
total_pages = pdf_document.page_count
last_page = pdf_document.load_page(total_pages - 1)
text = last_page.get_text()
pdf_document.close()
if "clauses" in text.lower():
return True
else:
return False
def Noc_invoice_pdf_to_img(pdf_path: str, folder_path: str, dpi: int = 300, quality: int = 95):
pdf_document = pymupdf.open(pdf_path)
folder_path = folder_path.rstrip(os.sep)
os.makedirs(folder_path, exist_ok=True)
pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]
total_pages = pdf_document.page_count
image_paths = []
for page_num in range(total_pages):
page = pdf_document.load_page(page_num)
pix = page.get_pixmap(dpi=dpi)
image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
output_path = os.path.join(folder_path, f"{pdf_name}_page_{page_num + 1}.jpg")
image.save(output_path, "JPEG", quality=quality)
image_paths.append(output_path)
pdf_document.close()
return image_paths
def delete_images(image_paths):
# Iterate through the list of image paths
for image_path in image_paths:
try:
# Check if the file exists before attempting to delete
if os.path.exists(image_path):
os.remove(image_path)
print(f"Deleted: {image_path}")
else:
print(f"File not found: {image_path}")
except Exception as e:
print(f"Error deleting {image_path}: {e}")
def noc_invoice_extraction(pdf_path: str, folder_path):
image_paths = Noc_invoice_pdf_to_img(pdf_path, folder_path)
data = {}
result = get_image_informations(image_paths[0], invoice_first_page_prompt, Noc_PurchaseOrder_information_parser)
data.update(result)
result = get_image_informations(image_paths[1], invoice_item_page1_prompt, Noc_PurchaseOrder_item1_parser)
data.update(result)
if Clauses_in_invoice(pdf_path):
for pic in range(len(image_paths) - 4):
new_item = get_image_informations(image_paths[pic + 2], invoice_item_pages_prompt, Noc_PurchaseOrder_items_parser)
for item in new_item["items"]:
data["items"].append(item)
result = get_image_informations(image_paths[-2], invoice_total_page_prompt, Noc_PurchaseOrder_total_parser)
data.update(result)
result = get_image_informations(image_paths[-1], invoice_clauses_page_prompt, Noc_PurchaseOrder_clauses_parser)
data.update(result)
delete_images(image_paths)
return data
else:
for pic in range(len(image_paths) - 3):
new_item = get_image_informations(image_paths[pic + 2], invoice_item_pages_prompt, Noc_PurchaseOrder_items_parser)
for item in new_item["items"]:
data["items"].append(item)
result = get_image_informations(image_paths[-2], invoice_total_page_prompt, Noc_PurchaseOrder_total_parser)
data.update(result)
delete_images(image_paths)
return data
def process_file(file, option):
if file is None:
return "Please upload a PDF or image file."
try:
save_dir = "uploaded_files"
os.makedirs(save_dir, exist_ok=True) # Create the directory if it doesn't exist
file_path = file.name
file_extension = os.path.splitext(file_path)[1].lower()
if file_extension in ['.pdf']:
# Process PDF files
if option == "Noc_timesheet_residential":
Noc_timeSheet_pdf_to_img(file_path, "output.jpg")
result = get_image_informations("output.jpg", Noc_Res_timesheet_prompt, Noc_Res_timeSheet_parser)
return result
elif option == "Noc_timesheet_rotational":
Noc_timeSheet_pdf_to_img(file_path, "output.jpg")
result = get_image_informations("output.jpg", Noc_Rot_timesheet_prompt, Noc_Rot_timeSheet_parser)
return result
elif option == "Noc_invoice":
result = noc_invoice_extraction(file_path, save_dir)
return result
elif file_extension in ['.jpg', '.jpeg', '.png']:
# Process image files directly
if option == "Noc_timesheet_residential":
result = get_image_informations(file_path, Noc_Res_timesheet_prompt, Noc_Res_timeSheet_parser)
return result
elif option == "Noc_timesheet_rotational":
result = get_image_informations(file_path, Noc_Rot_timesheet_prompt, Noc_Rot_timeSheet_parser)
return result
elif option == "Noc_invoice":
# For invoice images, we assume it's a single page
result = get_image_informations(file_path, invoice_first_page_prompt, Noc_PurchaseOrder_information_parser)
return result
else:
return "Unsupported file type. Please upload a PDF or image file."
except Exception as e:
return f"An error occurred: {e}"
# Define the Gradio interface
demo = gr.Interface(
fn=process_file,
inputs=[
gr.File(label="Upload PDF or Image"), # File upload input
gr.Radio(["Noc_timesheet_residential", "Noc_timesheet_rotational", "Noc_invoice"], label="Choose an option") # Radio buttons for options
],
outputs="text", # Text output
title="PDF/Image Processor",
description="Upload a PDF or image and choose an option to process the content."
)
with gr.Blocks() as app:
demo.render()
gr.Markdown("### PDF/Image examples") # Section title
with gr.Row():
gr.Image("TS.png", label="NOC timesheet example")
gr.Image("invoice.png", label="NOC invoice example")
app.launch()