update pdf library

app.py

@@ -3,56 +3,73 @@ import random
 import shutil
 import tempfile
 import zipfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 
+import fitz  # PyMuPDF
 import gradio as gr
-from huggingface_hub import HfApi, DatasetCard, DatasetCardData
-from pdf2image import convert_from_path
-from PyPDF2 import PdfReader
+from huggingface_hub import DatasetCard, DatasetCardData, HfApi
+
 from dataset_card_template import DATASET_CARD_TEMPLATE
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 
+def process_pdf(pdf_file, sample_size, temp_dir):
+    try:
+        pdf_path = pdf_file.name
+        doc = fitz.open(pdf_path)
+        total_pages = len(doc)
+
+        pages_to_convert = (
+            total_pages if sample_size == 0 else min(sample_size, total_pages)
+        )
+        selected_pages = (
+            sorted(random.sample(range(total_pages), pages_to_convert))
+            if sample_size > 0 and sample_size < total_pages
+            else range(total_pages)
+        )
+
+        images = []
+        for page_num in selected_pages:
+            page = doc[page_num]
+            pix = page.get_pixmap()
+            image_path = os.path.join(
+                temp_dir, f"{os.path.basename(pdf_path)}_page_{page_num+1}.png"
+            )
+            pix.save(image_path)
+            images.append(image_path)
+
+        doc.close()
+        return images, None
+    except Exception as e:
+        return [], f"Error processing {pdf_file.name}: {str(e)}"
+
+
 def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
     if not os.path.exists(temp_dir):
         os.makedirs(temp_dir)
+
     progress(0, desc="Starting conversion")
     all_images = []
     skipped_pdfs = []
-    for pdf_file in progress.tqdm(pdf_files, desc="Converting PDFs"):
-        try:
-            pdf_path = pdf_file.name
-            pdf = PdfReader(pdf_path)
-            total_pages = len(pdf.pages)
-
-            # Determine the number of pages to convert
-            pages_to_convert = (
-                total_pages if sample_size == 0 else min(sample_size, total_pages)
-            )
 
-            # Select random pages if sampling
-            if sample_size > 0 and sample_size < total_pages:
-                selected_pages = sorted(
-                    random.sample(range(1, total_pages + 1), pages_to_convert)
-                )
+    with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
+        future_to_pdf = {
+            executor.submit(process_pdf, pdf_file, sample_size, temp_dir): pdf_file
+            for pdf_file in pdf_files
+        }
+
+        for future in progress.tqdm(
+            as_completed(future_to_pdf), total=len(pdf_files), desc="Converting PDFs"
+        ):
+            pdf_file = future_to_pdf[future]
+            images, error = future.result()
+            if error:
+                skipped_pdfs.append(pdf_file.name)
+                gr.Info(error)
             else:
-                selected_pages = range(1, total_pages + 1)
-
-            # Convert selected PDF pages to images
-            for page_num in selected_pages:
-                images = convert_from_path(
-                    pdf_path, first_page=page_num, last_page=page_num
-                )
-                for image in images:
-                    image_path = os.path.join(
-                        temp_dir, f"{os.path.basename(pdf_path)}_page_{page_num}.jpg"
-                    )
-                    image.save(image_path, "JPEG")
-                    all_images.append(image_path)
-        except Exception as e:
-            skipped_pdfs.append(pdf_file.name)
-            gr.Info(f"Skipped PDF {pdf_file.name} due to error: {str(e)}")
+                all_images.extend(images)
 
     message = f"Saved {len(all_images)} images to temporary directory"
     if skipped_pdfs: