Update app.py

app.py

@@ -11,6 +11,11 @@ from PIL import Image
 import glob
 import pandas as pd
 import time
+from pdf2image import convert_from_path
+import pymupdf
+import camelot
+import numpy as np
+import fitz
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
@@ -34,6 +39,14 @@ def filter_detections(detections, target_class_name="mark"):
     detections.data['class_name'] = filtered_class_name
     return detections
 
+
+def add_label_detection(detections):
+    updated_class = [f"{class_name} {i + 1}" for i, class_name in enumerate(detections.data['class_name'])]
+    updated_id = [class_id + i for i, class_id in enumerate(detections.class_id)]
+    detections.data['class_name'] = np.array(updated_class)
+    detections.class_id = np.array(updated_id)
+    return detections
+
     
 def ends_with_number(s):
     return s[-1].isdigit()
@@ -185,6 +198,99 @@ def inference(
     return annotated_image, output_path_list
 
 
+def read_table(sheet):
+    excel_path = "output_tables.xlsx"
+    if os.path.exists(excel_path):
+        sheetnames = pd.ExcelFile(excel_path).sheet_names
+        if sheet in sheetnames:
+            df = pd.read_excel(excel_path, sheet_name=sheet)
+        else:
+            df = pd.DataFrame()
+    else:
+        df = pd.DataFrame()
+    return df
+
+
+def validate_df(df):
+    columns = []
+    count = 1
+    for col in df.columns:
+        if type(col) == int:
+            columns.append(f"Col {count}")
+            count += 1
+        else:
+            columns.append(col)
+    df.columns = columns
+    return df
+
+
+def analyze_table(file, conf_threshold, iou_threshold, progress=gr.Progress()):
+    progress(0, desc="Parsing table...")
+    img = convert_from_path(file)[0]
+    doc = pymupdf.open(file)
+    zoom_x = 1.0  # horizontal zoom
+    zoom_y = 1.0  # vertical zoom
+    mat = pymupdf.Matrix(zoom_x, zoom_y)
+
+    for i, page in enumerate(doc):
+        pix = page.get_pixmap(matrix=mat)
+        pix.save("temp.png")
+    image = cv2.imread("temp.png")
+    file_height, file_width, _ = image.shape
+    results = onnx_model_table(image, conf=conf_threshold, iou=iou_threshold, imgsz=640)[0]
+    detections = sv.Detections.from_ultralytics(results)
+    detections = add_label_detection(detections)
+    parsed_detections = parse_detection(detections)
+    # print(parsed_detections)
+    output_dir = "output_table"
+    # Check if the output directory exists, clear all the files inside
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    else:
+        for f in os.listdir(output_dir):
+            os.remove(os.path.join(output_dir, f))
+
+    box_annotator = sv.BoxAnnotator()
+    label_annotator = sv.LabelAnnotator(text_position=sv.Position.TOP_LEFT, text_thickness=1, text_padding=2)
+    annotated_image = image.copy()
+    annotated_image = box_annotator.annotate(
+        scene=annotated_image,
+        detections=detections
+    )
+    annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
+
+    pdf = fitz.open(file)
+    pdf_page = pdf[0]
+    table_area = [(ind,
+                   fitz.Rect(det['left'], det['top'], det['left'] + det['width'], det['top'] + det['height']))
+                  for ind, det in enumerate(parsed_detections)
+                  ]
+    table_list = []
+    for ind, area in progress.tqdm(table_area):
+
+        pdf_tabs = pdf_page.find_tables(clip=area)
+        if len(pdf_tabs.tables) > 0:
+            pdf_df = pdf_tabs[0].to_pandas()
+            print("Fitz Table Found!")
+        else:
+            cur = parsed_detections[ind]
+            table_areas = [f"{cur['left']},{file_height - cur['top']},{cur['left'] + cur['width']},{file_height - (cur['top'] + cur['height'])}"]
+            tables = camelot.read_pdf(file, pages='0', flavor='stream', row_tol=10, table_areas=table_areas)
+            pdf_df = tables[0].df
+            print("Camelot Table Found!")
+        pdf_df = validate_df(pdf_df)
+        table_list.append(pdf_df)
+    excel_path = "output_tables.xlsx"
+    sheet_list = []
+    with pd.ExcelWriter(excel_path, engine='xlsxwriter') as writer:
+        for i in range(len(table_list)):
+            sheet_name = f"Table_{i + 1}"
+            table_list[i].to_excel(writer, sheet_name=sheet_name, index=False)
+            sheet_list.append(sheet_name)
+
+    return img, annotated_image, excel_path, ", ".join(sheet_list)
+
+    
 TITLE = "<h1 style='font-size: 2.5em; text-align: center;'>Identify objects in construction design</h1>"
 DESCRIPTION = """<p style='font-size: 1.5em; line-height: 1.6em; text-align: left;'>Welcome to the object 
 identification application. This tool allows you to upload an image, and it will identify and annotate objects within 
@@ -211,6 +317,7 @@ EXAMPLES = [
 with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
     gr.HTML(TITLE)
     gr.HTML(DESCRIPTION)
+    
     with gr.Tab(label="Identify objects"):
         with gr.Row(equal_height=False):
             input_img = gr.Image(type="filepath", label="Upload Image")
@@ -239,5 +346,39 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
                     outputs=[output_img, gallery],
                     cache_examples=False,
                 )
-
+        
+    with gr.Tab(label="Detect and read table"):
+        with gr.Row():
+            with gr.Column():
+                upload_pdf = gr.Image(label="Upload PDF file")
+                upload_button = gr.UploadButton(label="Upload PDF file", file_types=[".pdf"])
+            with gr.Column():
+                output_img = gr.Image(label="Output Image", interactive=False)
+        
+        with gr.Row():
+            with gr.Column():
+                conf_thres_table = gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05,
+                                             label="Confidence Threshold")
+            with gr.Column():
+                iou_table = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="IOU Threshold")
+        
+        with gr.Row():
+            with gr.Column():
+                text_output = gr.Textbox(label="Table List")
+            with gr.Column():
+                file_output = gr.File()
+        
+        with gr.Row():
+            sheet_name = gr.Dropdown(choices=SHEET_LIST, allow_custom_value=True, label="Sheet Name")
+        
+        with gr.Row():
+            output_df = gr.Dataframe(label="Results")
+        upload_button.upload(analyze_table, [upload_button, conf_thres_table, iou_table],
+                             [upload_pdf, output_img, file_output, text_output])
+        conf_thres_table.change(analyze_table, [upload_button, conf_thres_table, iou_table],
+                                [upload_pdf, output_img, file_output, text_output])
+        iou_table.change(analyze_table, [upload_button, conf_thres_table, iou_table],
+                         [upload_pdf, output_img, file_output, text_output])
+        sheet_name.change(read_table, sheet_name, output_df)
+        
 demo.launch(debug=True)