pdfextract_fun.py

import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
# Filter warnings about inputs not requiring gradients
warnings.filterwarnings("ignore", message="None of the inputs have requires_grad=True. Gradients will be None")
warnings.filterwarnings("ignore", message="torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument.")

import cv2
import os
import fitz  # PyMuPDF
import numpy as np
import re
import pytesseract
import torch
from PIL import Image
from tqdm import tqdm 

from unilm.dit.object_detection.ditod import add_vit_config

from detectron2.config import CfgNode as CN
from detectron2.config import get_cfg
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor


# Step 1: instantiate config
cfg = get_cfg()
add_vit_config(cfg)
cfg.merge_from_file("cascade_dit_base.yml")

# Step 2: add model weights URL to config
cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth"

# Step 3: set device
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
#cfg.MODEL.DEVICE =  "cuda"

# Step 4: define model
predictor = DefaultPredictor(cfg)

def analyze_image(img):
    
    md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
    if cfg.DATASETS.TEST[0]=='icdar2019_test':
        md.set(thing_classes=["table"])
    else:
        md.set(thing_classes=["text","title","list","table","figure"])
    
    output = predictor(img)["instances"]
    v = Visualizer(img[:, :, ::-1],
                    md,
                    scale=1.0,
                    instance_mode=ColorMode.SEGMENTATION)
    result = v.draw_instance_predictions(output.to("cpu"))
    result_image = result.get_image()[:, :, ::-1]
    
    return result_image, output, v



def convert_pdf_to_jpg(pdf_path, output_folder, zoom_factor=2):
    doc = fitz.open(pdf_path)
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)

        # Adjust zoom factor for higher resolution
        mat = fitz.Matrix(zoom_factor, zoom_factor)  # Create a Matrix with the zoom factor
        pix = page.get_pixmap(matrix=mat)  # Render the page using the matrix

        output_file = f"{output_folder}/page_{page_num}.jpg"
        pix.save(output_file)



def process_jpeg_images(output_folder):
     for page_num in tqdm(range(len(os.listdir(output_folder))), desc="Processing the pdf"):
         file_path = f"{output_folder}/page_{page_num}.jpg"
         img = cv2.imread(file_path)
         if img is None:
             print(f"Failed to read {file_path}. Skipping.")
             continue
         result_image, output, v = analyze_image(img)

         # Saving logic
         save_extracted_instances(img, output, page_num,output_folder)



def save_extracted_instances(img, output, page_num, dest_folder, confidence_threshold=0.8):
    class_names = {
        0: "text",
        1: "title",
        2: "list",
        3: "table",
        4: "figure"
    }

    threshold_value = 0  # Standard deviation threshold
    min_height = 0  # Minimum height threshold

    instances = output.to("cpu")
    boxes = instances.pred_boxes.tensor.numpy()
    class_ids = instances.pred_classes.tolist()
    scores = instances.scores.tolist()  # Get prediction scores

    image_counter = 1
    for box, class_id, score in zip(boxes, class_ids, scores):
        # Check if the prediction score meets the confidence threshold
        if score >= confidence_threshold:
            class_name = class_names.get(class_id, "unknown")

            # Save only if class is 'figure' or 'table'
            if class_name in ["figure", "table","text"]:
                x1, y1, x2, y2 = map(int, box)
                cropped_image = img[y1:y2, x1:x2]

                if np.std(cropped_image) > threshold_value and (y2 - y1) > min_height:
                    save_path = os.path.join(dest_folder, f"page_{page_num}_{class_name}_{image_counter}.jpg")
                    cv2.imwrite(save_path, cropped_image)
                    image_counter += 1


def delete_files_in_folder(folder_path):
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        if os.path.isfile(file_path):
            os.remove(file_path)



def rename_files_sequentially(folder_path):
    # Regex pattern to match 'page_{page_num}_{class_name}_{image_counter}.jpg'
    pattern = re.compile(r'page_(\d+)_(\w+)_(\d+).jpg', re.IGNORECASE)

    # List files in the folder
    files = os.listdir(folder_path)

    # Filter and sort files based on the regex pattern
    sorted_files = sorted(
        [f for f in files if pattern.match(f)],
        key=lambda x: (int(pattern.match(x).group(1)), pattern.match(x).group(2).lower(), int(pattern.match(x).group(3)))
    )

    # Initialize an empty dictionary for counters
    counters = {}

    for filename in sorted_files:
        match = pattern.match(filename)
        if match:
            page_num, class_name, _ = match.groups()
            class_name = class_name.lower()  # Convert class name to lowercase

            # Initialize counter for this class if it doesn't exist
            if class_name not in counters:
                counters[class_name] = 1

            # New filename format: '{class_name}_{sequential_number}.jpg'
            new_filename = f"{class_name}_{counters[class_name]}.jpg"
            counters[class_name] += 1

            # Rename the file
            os.rename(os.path.join(folder_path, filename), os.path.join(folder_path, new_filename))

            #print(f"Renamed '{filename}' to '{new_filename}'")


def ocr_folder(folder_path):
    # Regex pattern to match 'text_{number}.jpg'
    pattern = re.compile(r'text_\d+\.jpg', re.IGNORECASE)

    # Create a subfolder for the OCR text files
    ocr_text_folder = os.path.join(folder_path, "ocr_results")
    if not os.path.exists(ocr_text_folder):
        os.makedirs(ocr_text_folder)

    for filename in os.listdir(folder_path):
        if pattern.match(filename):
            image_path = os.path.join(folder_path, filename)
            text = ocr_image(image_path)
            
            # Save the OCR result to a text file in the subfolder
            text_file_name = filename.replace('.jpg', '.txt')
            text_file_path = os.path.join(ocr_text_folder, text_file_name)
            with open(text_file_path, 'w') as file:
                file.write(text)
            
            #print(f"OCR result for {filename} saved to {text_file_path}\n")

def ocr_image(image_path):
    image = Image.open(image_path)
    text = pytesseract.image_to_string(image)
    return text