process in parallel

app.py

@@ -1,3 +1,4 @@
+import multiprocessing
 import os
 import random
 import shutil
@@ -9,13 +10,18 @@ from datetime import datetime
 import fitz  # PyMuPDF
 import gradio as gr
 from huggingface_hub import DatasetCard, DatasetCardData, HfApi
+from PIL import Image
 
 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, progress=gr.Progress()):
+CPU_COUNT = multiprocessing.cpu_count()
+MAX_WORKERS = min(32, CPU_COUNT)  # Use CPU count directly for processes
+
+
+def process_pdf(pdf_file, sample_size, temp_dir):
     try:
         pdf_path = pdf_file.name
         doc = fitz.open(pdf_path)
@@ -31,21 +37,20 @@ def process_pdf(pdf_file, sample_size, temp_dir, progress=gr.Progress()):
         )
 
         images = []
-        for page_num in progress.tqdm(
-            selected_pages, desc=f"Converting {os.path.basename(pdf_path)}", unit="page"
-        ):
+        for page_num in selected_pages:
             page = doc[page_num]
-            pix = page.get_pixmap()
+            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # Increase resolution
+            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
             image_path = os.path.join(
-                temp_dir, f"{os.path.basename(pdf_path)}_page_{page_num+1}.png"
+                temp_dir, f"{os.path.basename(pdf_path)}_page_{page_num+1}.jpg"
             )
-            pix.save(image_path)
+            image.save(image_path, "JPEG", quality=85, optimize=True)
             images.append(image_path)
 
         doc.close()
-        return images, None
+        return images, None, len(images)
     except Exception as e:
-        return [], f"Error processing {pdf_file.name}: {str(e)}"
+        return [], f"Error processing {pdf_file.name}: {str(e)}", 0
 
 
 def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
@@ -56,15 +61,26 @@ def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
     all_images = []
     skipped_pdfs = []
 
-    for i, pdf_file in enumerate(
-        progress.tqdm(pdf_files, desc="Converting PDFs", unit="PDF")
-    ):
-        images, error = process_pdf(pdf_file, sample_size, temp_dir, progress)
-        if error:
-            skipped_pdfs.append(pdf_file.name)
-            gr.Info(error)
-        else:
-            all_images.extend(images)
+    total_pages = sum(len(fitz.open(pdf.name)) for pdf in pdf_files)
+    processed_pages = 0
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+        future_to_pdf = {
+            executor.submit(process_pdf, pdf, sample_size, temp_dir): pdf
+            for pdf in pdf_files
+        }
+
+        for future in as_completed(future_to_pdf):
+            pdf = future_to_pdf[future]
+            images, error, pages_processed = future.result()
+            if error:
+                skipped_pdfs.append(error)
+                gr.Info(error)
+            else:
+                all_images.extend(images)
+
+            processed_pages += pages_processed
+            progress((processed_pages / total_pages), desc=f"Processing {pdf.name}")
 
     message = f"Saved {len(all_images)} images to temporary directory"
     if skipped_pdfs:
@@ -138,7 +154,7 @@ def process_pdfs(
                     repo_type="dataset",
                     private=private_repo,
                 )
-                hf_api.upload_large_folder(
+                hf_api.upload_folder(
                     folder_path=temp_dir,
                     repo_id=hf_repo,
                     repo_type="dataset",