917eedffcf
Main / Python 3.11 - Docs (push) Waiting to run
Main / Python 3.11 - Build (push) Waiting to run
Main / Python 3.11 - Lint (push) Waiting to run
Main / Python 3.11 - Style (push) Waiting to run
Main / Python 3.11 - Test (push) Waiting to run
Main / GPU CI (push) Blocked by required conditions
Main / Release (push) Blocked by required conditions
Main / Build and Push Docker Images (push) Blocked by required conditions
111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Sample random pages from PDFs in an input directory and save them as individual PDFs.
|
|
|
|
This script takes PDFs from an input directory, randomly selects N pages from each,
|
|
and writes each selected page as a separate PDF file in the output directory.
|
|
Uses a fixed random seed for reproducibility.
|
|
"""
|
|
|
|
import argparse
|
|
import random
|
|
from pathlib import Path
|
|
|
|
from pypdf import PdfReader, PdfWriter
|
|
|
|
|
|
def sample_pdf_pages(input_dir: Path, output_dir: Path, num_pages: int, seed: int = 42):
|
|
"""
|
|
Sample random pages from PDFs and save them as individual files.
|
|
|
|
Args:
|
|
input_dir: Directory containing input PDF files
|
|
output_dir: Directory to save sampled pages
|
|
num_pages: Number of pages to sample from each PDF
|
|
seed: Random seed for reproducibility
|
|
"""
|
|
# Set random seed for consistency
|
|
random.seed(seed)
|
|
|
|
# Create output directory if it doesn't exist
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Get all PDF files from input directory
|
|
pdf_files = sorted(input_dir.glob("*.pdf"))
|
|
|
|
if not pdf_files:
|
|
print(f"No PDF files found in {input_dir}")
|
|
return
|
|
|
|
print(f"Found {len(pdf_files)} PDF file(s)")
|
|
|
|
for pdf_path in pdf_files:
|
|
print(f"\nProcessing: {pdf_path.name}")
|
|
|
|
try:
|
|
reader = PdfReader(pdf_path)
|
|
total_pages = len(reader.pages)
|
|
|
|
print(f" Total pages: {total_pages}")
|
|
|
|
# Determine how many pages to sample (min of requested and available)
|
|
pages_to_sample = min(num_pages, total_pages)
|
|
|
|
if pages_to_sample < num_pages:
|
|
print(f" Warning: Only {total_pages} pages available, sampling all")
|
|
|
|
# Randomly select page indices
|
|
selected_indices = sorted(random.sample(range(total_pages), pages_to_sample))
|
|
print(f" Selected page indices: {selected_indices}")
|
|
|
|
# Extract and save each selected page
|
|
for page_idx in selected_indices:
|
|
# Create writer for single page
|
|
writer = PdfWriter()
|
|
writer.add_page(reader.pages[page_idx])
|
|
|
|
# Create output filename: original_name_page_N.pdf
|
|
output_filename = f"{pdf_path.stem}_page_{page_idx}.pdf"
|
|
output_path = output_dir / output_filename
|
|
|
|
# Write the single page PDF
|
|
with open(output_path, "wb") as output_file:
|
|
writer.write(output_file)
|
|
|
|
print(f" Saved: {output_filename}")
|
|
|
|
except Exception as e:
|
|
print(f" Error processing {pdf_path.name}: {e}")
|
|
continue
|
|
|
|
print(f"\nDone! Output saved to {output_dir}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Sample random pages from PDFs and save as individual files")
|
|
parser.add_argument("--input-dir", type=Path, required=True, help="Directory containing input PDF files")
|
|
parser.add_argument("--output-dir", type=Path, required=True, help="Directory to save sampled pages")
|
|
parser.add_argument("--pages", type=int, required=True, help="Number of pages to sample from each PDF")
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility (default: 42)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Validate arguments
|
|
if not args.input_dir.exists():
|
|
print(f"Error: Input directory does not exist: {args.input_dir}")
|
|
return
|
|
|
|
if not args.input_dir.is_dir():
|
|
print(f"Error: Input path is not a directory: {args.input_dir}")
|
|
return
|
|
|
|
if args.pages < 1:
|
|
print(f"Error: Number of pages must be at least 1")
|
|
return
|
|
|
|
sample_pdf_pages(args.input_dir, args.output_dir, args.pages, args.seed)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|