add progress

app.py

@@ -15,7 +15,7 @@ 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):
+def process_pdf(pdf_file, sample_size, temp_dir, progress=gr.Progress()):
     try:
         pdf_path = pdf_file.name
         doc = fitz.open(pdf_path)
@@ -31,7 +31,7 @@ def process_pdf(pdf_file, sample_size, temp_dir):
         )
 
         images = []
-        for page_num in selected_pages:
+        for i, page_num in enumerate(selected_pages):
             page = doc[page_num]
             pix = page.get_pixmap()
             image_path = os.path.join(
@@ -39,6 +39,10 @@ def process_pdf(pdf_file, sample_size, temp_dir):
             )
             pix.save(image_path)
             images.append(image_path)
+            progress(
+                (i + 1) / len(selected_pages),
+                desc=f"Converting {os.path.basename(pdf_path)}",
+            )
 
         doc.close()
         return images, None
@@ -56,13 +60,13 @@ def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
 
     with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
         future_to_pdf = {
-            executor.submit(process_pdf, pdf_file, sample_size, temp_dir): pdf_file
+            executor.submit(
+                process_pdf, pdf_file, sample_size, temp_dir, progress
+            ): 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"
-        ):
+        for i, future in enumerate(as_completed(future_to_pdf)):
             pdf_file = future_to_pdf[future]
             images, error = future.result()
             if error:
@@ -70,6 +74,7 @@ def pdf_to_images(pdf_files, sample_size, temp_dir, progress=gr.Progress()):
                 gr.Info(error)
             else:
                 all_images.extend(images)
+            progress((i + 1) / len(pdf_files), desc="Converting PDFs")
 
     message = f"Saved {len(all_images)} images to temporary directory"
     if skipped_pdfs: