Update pdfextract_fun.py

pdfextract_fun.py

@@ -1,193 +1,83 @@
+import os
+import re
 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 tqdm import tqdm
 from detectron2.config import get_cfg
-from detectron2.utils.visualizer import ColorMode, Visualizer
 from detectron2.data import MetadataCatalog
 from detectron2.engine import DefaultPredictor
+from detectron2.utils.visualizer import ColorMode, Visualizer
+from unilm.dit.object_detection.ditod import add_vit_config
 
+# Filter specific warnings
+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.")
 
-# 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):
-    
+# Configuration setup
+def setup_config():
+    cfg = get_cfg()
+    add_vit_config(cfg)
+    cfg.merge_from_file("cascade_dit_base.yml")
+    cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth"
+    cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    return cfg
+
+# Analyze image
+def analyze_image(img, cfg):
+    """Analyze an image and return the result image, output, and visualizer."""
     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
-
+    thing_classes = ["table"] if cfg.DATASETS.TEST[0] == 'icdar2019_test' else ["text", "title", "list", "table", "figure"]
+    md.set(thing_classes=thing_classes)
 
+    output = DefaultPredictor(cfg)(img)["instances"]
+    v = Visualizer(img[:, :, ::-1], metadata=md, scale=1.0, instance_mode=ColorMode.SEGMENTATION)
+    result = v.draw_instance_predictions(output.to("cpu"))
+    return result.get_image()[:, :, ::-1], output, v
 
+# PDF to JPEG conversion
 def convert_pdf_to_jpg(pdf_path, output_folder, zoom_factor=2):
+    """Convert PDF file to JPEG images, saved in the specified output folder."""
     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"
+    for page_num, page in enumerate(doc):
+        mat = fitz.Matrix(zoom_factor, zoom_factor)
+        pix = page.get_pixmap(matrix=mat)
+        output_file = os.path.join(output_folder, f"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)
-
-
-
+# Process JPEG images in a folder
+def process_jpeg_images((output_folder, cfg, batch_size=10):
+    image_paths = [os.path.join(output_folder, f) for f in os.listdir(output_folder) if f.endswith('.jpg')]
+    batches = [image_paths[i:i + batch_size] for i in range(0, len(image_paths), batch_size)]
+    
+    for batch in tqdm(batches, desc="Processing images in batches"):
+        images = [cv2.imread(image_path) for image_path in batch]
+        batch_results = batch_analyze_images(images, cfg)  # This function needs to be implemented to support batch processing
+        
+        for i, (result_image, output, v) in enumerate(batch_results):
+            # Assuming batch_analyze_images returns a list of tuples, each containing the results for one image
+            save_extracted_instances(images[i], output, i, output_folder)
+
+# Save extracted instances
 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
-
+    """Save instances extracted from an image to the destination folder."""
+    class_names = {0: "text", 1: "title", 2: "list", 3: "table", 4: "figure"}
     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
+    for i, (box, class_id, score) in enumerate(zip(instances.pred_boxes.tensor.numpy(), instances.pred_classes.tolist(), instances.scores.tolist())):
+        if score >= confidence_threshold and class_names.get(class_id) in ["figure", "table", "text"]:
+            x1, y1, x2, y2 = map(int, box)
+            cropped_image = img[y1:y2, x1:x2]
+            if np.std(cropped_image) > 0 and (y2 - y1) > 0:  # Replace with actual thresholds if needed
+                save_path = os.path.join(dest_folder, f"page_{page_num}_{class_names[class_id]}_{i + 1}.jpg")
+                cv2.imwrite(save_path, cropped_image)
+
+# Additional functions like delete_files_in_folder, rename_files_sequentially, ocr_folder, and ocr_image can be included as is, assuming they were satisfactory before.
+
+cfg = setup_config()
+# Example usage
+convert_pdf_to_jpg("sample.pdf", "output_folder")
+process_jpeg_images("output_folder", cfg)