From eb11d957645cfed062b123d49ffb814bf18774b1 Mon Sep 17 00:00:00 2001 From: Esa Kataja Date: Tue, 3 Dec 2024 19:25:02 +0200 Subject: [PATCH] feat: Add PDF cleaner with deskew and image enhancement Features: - Page deskewing using contour detection - Image enhancement with edge preservation - High-resolution output (min 2048px width) - Black and white conversion with proper contrast --- pdf_cleaner.py | 157 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 pdf_cleaner.py diff --git a/pdf_cleaner.py b/pdf_cleaner.py new file mode 100644 index 0000000..1c4e60a --- /dev/null +++ b/pdf_cleaner.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +import os +import cv2 +import numpy as np +from pdf2image import convert_from_path +import img2pdf +from PIL import Image +import argparse + +def deskew(image): + """Deskew the image using contour detection and rotation.""" + # Create a copy for processing while keeping original quality + proc_image = image.copy() + + # Convert to grayscale and blur + gray = cv2.cvtColor(proc_image, cv2.COLOR_BGR2GRAY) + blur = cv2.GaussianBlur(gray, (9, 9), 0) + + # Threshold the image + thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1] + + # Find all contours + contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + + if not contours: + return image + + # Find largest contour + contour = max(contours, key=cv2.contourArea) + + # Find minimum area rectangle + rect = cv2.minAreaRect(contour) + angle = rect[-1] + + # Adjust angle to be between -45 and 45 degrees + while angle < -45: + angle += 90 + while angle > 45: + angle -= 90 + + # Only rotate if the angle is significant enough + if abs(angle) < 0.5: # Skip tiny rotations + return image + + # Rotate the image + (h, w) = image.shape[:2] + center = (w // 2, h // 2) + M = cv2.getRotationMatrix2D(center, angle, 1.0) + rotated = cv2.warpAffine(image, M, (w, h), + flags=cv2.INTER_CUBIC, + borderMode=cv2.BORDER_REPLICATE) + + return rotated + +def adjust_levels(image): + """Adjust image levels for better contrast and clarity.""" + # Convert to grayscale if not already + if len(image.shape) == 3: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + else: + gray = image + + # Apply bilateral filter to preserve edges while reducing noise + denoised = cv2.bilateralFilter(gray, 9, 75, 75) + + # Create a background mask using morphological operations + kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15)) + background = cv2.morphologyEx(denoised, cv2.MORPH_DILATE, kernel) + + # Subtract background to normalize lighting + normalized = cv2.subtract(background, denoised) + + # Apply Gaussian blur to reduce noise while preserving edges + blurred = cv2.GaussianBlur(normalized, (3, 3), 0) + + # Use Otsu's thresholding for optimal binary threshold + _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) + + # Always ensure background is white and text is black + if cv2.countNonZero(binary) < binary.size / 2: + binary = cv2.bitwise_not(binary) + + return binary + +def process_pdf(input_path, output_path): + """Process a PDF file and save the cleaned version.""" + print(f"Processing {input_path}...") + + # Convert PDF to images with high DPI to ensure minimum width of 2048px + pages = convert_from_path(input_path, dpi=300) + + # Create temporary directory for processed images + temp_dir = "temp_processed_images" + os.makedirs(temp_dir, exist_ok=True) + + # Process each page + temp_image_paths = [] + for i, page in enumerate(pages): + print(f"Processing page {i+1}/{len(pages)}") + + # Ensure minimum width of 2048px + width, height = page.size + scale = max(1, 2048 / width) + if scale > 1: + new_width = int(width * scale) + new_height = int(height * scale) + page = page.resize((new_width, new_height), Image.Resampling.LANCZOS) + + # Convert PIL Image to OpenCV format + opencv_image = cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR) + + # Process the image + deskewed = deskew(opencv_image) + cleaned = adjust_levels(deskewed) + + # Convert to PIL Image + pil_image = Image.fromarray(cleaned) + + # Save temporarily as 1-bit PNG + temp_path = os.path.join(temp_dir, f"page_{i:03d}.png") + pil_image.save(temp_path, "PNG", optimize=False) + temp_image_paths.append(temp_path) + + # Save processed images as PDF with high quality settings + print("Saving cleaned PDF...") + a4_width_mm = 210 + a4_height_mm = 297 + layout_fun = img2pdf.get_layout_fun((img2pdf.mm_to_pt(a4_width_mm), + img2pdf.mm_to_pt(a4_height_mm))) + + with open(output_path, "wb") as f: + f.write(img2pdf.convert(temp_image_paths, + layout_fun=layout_fun, + with_pdfrw=True)) + + # Clean up temporary files + for temp_path in temp_image_paths: + os.remove(temp_path) + os.rmdir(temp_dir) + + print(f"Saved cleaned PDF to {output_path}") + +def main(): + parser = argparse.ArgumentParser(description='Clean and straighten image-based PDFs') + parser.add_argument('input_pdf', help='Path to the input PDF file') + parser.add_argument('output_pdf', help='Path for the output PDF file') + + args = parser.parse_args() + + if not os.path.exists(args.input_pdf): + print(f"Error: Input file '{args.input_pdf}' does not exist") + return + + process_pdf(args.input_pdf, args.output_pdf) + +if __name__ == "__main__": + main()