utils.py

import subprocess
import streamlit as st
import cv2
import numpy as np
from PIL import Image
import pytesseract
def get_pdf_page_count(pdf_path):
    try:
        # Running pdfinfo command to get information about the PDF
        result = subprocess.run(['pdfinfo', pdf_path], stdout=subprocess.PIPE, text=True)
        # Parsing the output to find the line with the number of pages
        for line in result.stdout.split('\n'):
            if 'Pages:' in line:
                return int(line.split(':')[1].strip())
    except Exception as e:
        print(f"An error occurred: {e}")
        return None
#configurable extract rectange rectangle size

def extract_rectangle_from_image(gray, min_width, min_height):
    bounding_boxes = []
    #gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    #edges = cv2.Canny(gray, 10, 200, apertureSize=3)
    kernel = np.ones((3,3), np.uint8)
    dilated_edges = cv2.dilate(edges, kernel, iterations=1)
    contours, _ = cv2.findContours(dilated_edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    #contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)


    index = 0
    for cnt in contours:
        approx = cv2.approxPolyDP(cnt, 0.01*cv2.arcLength(cnt, True), True)
        #approx = cv2.approxPolyDP(cnt, 0.1*cv2.arcLength(cnt, True), True)
        if len(approx) == 4:  # Rectangle check
            x, y, w, h = cv2.boundingRect(approx)
            # print(f"x: {x}, y: {y}, w: {w}, h: {h}")
            if w >= min_width and h >= min_height:
                bounding_boxes.append((x, y, w, h))
                #print(x, y, w, h)
    return bounding_boxes
def is_close(box1, box2, threshold=10):
    # Calculate the distance between the top-left corners of the two boxes
    distance = ((box1[0] - box2[0]) ** 2 + (box1[1] - box2[1]) ** 2) ** 0.5
    return distance < threshold
def remove_close_boxes(boxes, threshold=10):
    kept_boxes = []
    for box in boxes:
        # Assume the box is not close to others by default
        is_close_to_others = False
        for kept_box in kept_boxes:
            if is_close(box, kept_box, threshold):
                is_close_to_others = True
                break
        # If the box is not close to any box we've kept, add it to the list of kept boxes
        if not is_close_to_others:
            kept_boxes.append(box)
    return kept_boxes
def is_contained(box1, box2):
    """
    Check if box1 is contained within box2.
    Each box is defined as (x, y, w, h).
    """
    x1, y1, w1, h1 = box1
    x2, y2, w2, h2 = box2

    # Check if all corners of box1 are inside box2
    return x2 <= x1 and y2 <= y1 and x2 + w2 >= x1 + w1 and y2 + h2 >= y1 + h1

def remove_contained_boxes(boxes):
    """
    Remove boxes that are contained within other boxes.
    """
    non_contained_boxes = []

    for i, box1 in enumerate(boxes):
        # Check if there's another box that contains box1
        if not any(is_contained(box1, box2) for j, box2 in enumerate(boxes) if i != j):
            non_contained_boxes.append(box1)

    return non_contained_boxes
def draw_colored_boxes_on_image_np(image, boxes_list,color_tuple):
  for x, y, w, h in boxes_list:
      #x, y, w, h = box[0]
      cv2.rectangle(image, (x, y), (x + w, y + h), color_tuple, thickness=5)

def is_filled_rectangle(image, rect, background_threshold=10, variance_threshold=0.1):

    x, y, w, h = rect
    roi = image[y+1:y+h-1, x+1:x+w-1]

    return np.all(roi == 0)
def get_below_box(image_np, x, y,width,step=15):
  #print("x,y,width="+str(x)+","+str(y)+","+str(width))

  index_y = -1
  #print("get_below_box"+str(image_np.shape))
  if y+step < image_np.shape[0]:
    index_y = y
    while index_y+step < image_np.shape[0]:
      #print(str( np.all(image_np[index_y:index_y+step,x:x+width] == 255)))

      # image_np_copy = image_np.copy()
      # bgr_image = cv2.cvtColor(image_np_copy, cv2.COLOR_GRAY2BGR)
      # cv2.rectangle(bgr_image, (x, index_y), (x + width, index_y +step), color_tuple, thickness=5)
      # display_image_np(bgr_image)


      if np.all(image_np[index_y:index_y+step,x:x+width] == 255):
        # index_y += step
        break
      index_y += step
  return index_y
def get_above_box(image_np, x, y,width,step=15):
  #print("x,y,width="+str(x)+","+str(y)+","+str(width))

  index_y = -1
  #print("get_below_box"+str(image_np.shape))
  if y-step > 0:
    index_y = y
    while index_y-step > 0:
      #print(str( np.all(image_np[index_y:index_y+step,x:x+width] == 255)))

      # image_np_copy = image_np.copy()
      # bgr_image = cv2.cvtColor(image_np_copy, cv2.COLOR_GRAY2BGR)
      # color_tuple=(0, 255, 0)
      # cv2.rectangle(bgr_image, (x, index_y-step), (x + width, index_y), color_tuple, thickness=5)
      # display_image_np(bgr_image)


      if np.all(image_np[index_y-step:index_y,x:x+width] == 255):
        # index_y += step
        break
      index_y -= step
  return index_y
def is_note_rectangle(image_np, rect):
  x, y, w, h = rect
  roi = image_np[y+1:y+h-1, x+1:x+w-1]
  roi_converted = Image.fromarray(cv2.cvtColor(roi, cv2.COLOR_BGR2RGB))
  text = pytesseract.image_to_string(roi_converted)
  text = text.strip()
  note_str="note"
  print("is note text box="+str(text.lower().startswith(note_str.lower())))
  return text.lower().startswith(note_str.lower())
def extract_bounding_boxes_from_image_np(image_np, bounding_boxes_list, above_check_offset, above_caption_offset, color_tuple):

  image_np_copy=image_np.copy()
  rect_content_list=[]
  above_rect_content_list=[]
  figures_image_list=[]
  tables_image_list=[]
  index = 0
  for box in bounding_boxes_list:
    x, y, w, h = box
    if not is_filled_rectangle(image_np_copy, box):
      # print("box="+str(box)+"not filled")
      y_index= get_below_box(image_np, x, y+h,w)
      if y_index == -1 or is_note_rectangle(image_np_copy, box):
        # print("below text not found")
        rect_content =image_np[y:y+h, x:x+w]
        # rect_content_list.append(rect_content)
        cv2.rectangle(image_np_copy, (x, y), (x+w, y+h), color_tuple, cv2.FILLED)
      else:
        # print("below text found")
        rect_content =image_np[y:y_index, x:x+w]
        # rect_content_list.append(rect_content)
        cv2.rectangle(image_np_copy, (x, y), (x+w, y_index), color_tuple, cv2.FILLED)

      cv2.rectangle(image_np_copy, (x, y), (x+w, y+h), color_tuple, cv2.FILLED)

      above_box_y= get_above_box(image_np, x, y,w)
      if above_box_y == -1 or above_box_y == y:
        # print("box="+str(box)+"no above box")
        above_rect_content_list.append(None)
        rect_content_list.append(rect_content)
      else:
        # print("box="+str(box)+"above box exist")
        above_rect_content = image_np[above_box_y:y, x:x+w]
        # above_rect_content_list.append(above_rect_content)
        above_converted = Image.fromarray(cv2.cvtColor(above_rect_content, cv2.COLOR_BGR2RGB))
        text = pytesseract.image_to_string(above_converted)
        text = text.strip()
        figure_str ="Figure"
        table_str ="Table"
        if text.lower().startswith(figure_str.lower()):
          print(text)
          figures_image_list.append((text,rect_content))

        elif  text.lower().startswith(table_str.lower()):
          print(text)
          tables_image_list.append((text,rect_content))
        else:
          above_rect_content_list.append((text, rect_content))
        rect_content_list.append(rect_content)

        cv2.rectangle(image_np_copy, (x, above_box_y), (x+w, y), color_tuple, cv2.FILLED)
      # above_rect_content = image_np[y-above_check_offset:y, x:x+w]
      # if np.all(above_rect_content == 255):
      #   # print("box="+str(box)+"above all white")
      #   above_rect_content_list.append(None)
      # else:
      #   # print("box="+str(box)+"above not all white")
      #   above_rect_content = image_np[y-above_caption_offset:y, x:x+w]
      #   above_rect_content_list.append(above_rect_content)
      #   cv2.rectangle(image_np_copy, (x, y), (x+w, y-above_caption_offset), color_tuple, cv2.FILLED)

      index += 1
    # else:
    #   print("box="+str(box)+"filled")
  return rect_content_list,above_rect_content_list,  figures_image_list, tables_image_list, image_np_copy
def gray_pdf_image_np_to_text(image_index,gray_pdf_image_np, debug=False):
    bounding_boxes_list = extract_rectangle_from_image(gray_pdf_image_np, 500, 20)
    bounding_boxes_list = remove_close_boxes (bounding_boxes_list, 10)
    bounding_boxes_list = remove_contained_boxes(bounding_boxes_list)
    if debug:
        bgr_image = cv2.cvtColor(gray_pdf_image_np, cv2.COLOR_GRAY2BGR)
        color_tuple = (0, 255, 0)
        draw_colored_boxes_on_image_np(bgr_image, bounding_boxes_list, color_tuple)
        # st.image(Image.fromarray(bgr_image)) #to_be_displayed

    text_box_list, above_test_box_list,figures_image_list,tables_image_list, cropped_image = extract_bounding_boxes_from_image_np(gray_pdf_image_np,
                                                                                             bounding_boxes_list, 30,
                                                                                             50, (255, 255, 255))
    if debug:
        debug_text_box_index = 0
        for text_box, above_text_box in zip(text_box_list, above_test_box_list):
            print("text box start")
            if above_text_box is not None:
                st.write(above_text_box[0])
                st.image(Image.fromarray(above_text_box[1]))
                # st.write(text)
            st.image(Image.fromarray(text_box))
            debug_text_box_index = debug_text_box_index + 1
        for figure in figures_image_list:
            st.write(figure[0])
            st.image(Image.fromarray(figure[1]))
        for table in tables_image_list:
            st.write(table[0])
            st.image(Image.fromarray(table[1]))