Update app.py

app.py

@@ -6,9 +6,7 @@ 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):
+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
@@ -17,16 +15,15 @@ def Noc_timeSheet_pdf_to_img(pdf_path,output_path,dpi: int = 300, quality: int =
     # 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)
+    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))
+    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))
@@ -38,20 +35,8 @@ def Noc_timeSheet_pdf_to_img(pdf_path,output_path,dpi: int = 300, quality: int =
     combined_image.paste(croped2, (0, upper_height))
 
     # Save the combined image
-    combined_image.save(output_path, "JPEG",quality=quality)  
-    
-    #-----------S3------------ need S3_BUCKET,S3_REGION,S3_URL 
-    # import boto3
-    
-    # s3_client = boto3.client('s3', region_name=S3_REGION)
-    # s3_client.upload_file(output_path, S3_BUCKET, key)
-    
-    # file_url = f"{S3_URL}/{key}"
-    
-    # return file_url
-    
-    # return output_path
- 
+    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.
@@ -65,16 +50,15 @@ def Clauses_in_invoice(pdf_path: str) -> bool:
         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=[]
+    image_paths = []
     for page_num in range(total_pages):
         page = pdf_document.load_page(page_num)
         pix = page.get_pixmap(dpi=dpi)
@@ -82,19 +66,6 @@ def Noc_invoice_pdf_to_img(pdf_path: str, folder_path: str, dpi: int = 300, qual
 
         output_path = os.path.join(folder_path, f"{pdf_name}_page_{page_num + 1}.jpg")
         image.save(output_path, "JPEG", quality=quality)
-        
-        #-----------S3------------ need S3_BUCKET,S3_REGION,S3_URL 
-        # import boto3
-        
-        # s3_client = boto3.client('s3', region_name=S3_REGION)
-        # s3_client.upload_file(output_path, S3_BUCKET, key)
-        
-        # file_url = f"{S3_URL}/{key}"
-        
-        # append the s3 links
-        # image_paths.append(file_url)
-        
-        
         image_paths.append(output_path)
 
     pdf_document.close()
@@ -113,84 +84,92 @@ def delete_images(image_paths):
         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)
+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)
+    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)
+    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 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)
+        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)
+        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 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)
+        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_pdf(file, option):
+def process_file(file, option):
     if file is None:
-        return "Please upload a PDF file."
+        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
 
-        
-        # Save the uploaded file to the new location
         file_path = file.name
-        
-        # Process based on the selected option
-        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
-            
-    #     else:
-    #         return "Invalid option selected."
+        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_pdf,
+    fn=process_file,
     inputs=[
-        gr.File(label="Upload PDF"),  # File upload input
-        gr.Radio(["Noc_timesheet_residential","Noc_timesheet_rotational", "Noc_invoice"], label="Choose an option")  # Radio buttons for options
+        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 Processor",
-    description="Upload a PDF and choose an option to process the text."
+    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 examples")  # Section title
+    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()
+app.launch()