Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import tempfile | |
| import numpy as np | |
| import cv2 | |
| import img2pdf | |
| from pathlib import Path | |
| from pdf2image import convert_from_path # Make sure pdf2image is installed: pip install pdf2image opencv-python numpy img2pdf | |
| # Define a default bounding box (e.g., for a logo area) | |
| # Bounding box: (x1, y1, x2, y2) in fractions of image size | |
| default_bbox = (0.75, 0.9, 0.98, 0.98) | |
| def expand_bbox(bbox, expand_factor=0.3): | |
| """Expands a bounding box by a given factor.""" | |
| x_min, y_min, x_max, y_max = bbox | |
| width = x_max - x_min | |
| height = y_max - y_min | |
| new_x_min = max(0, x_min - width * expand_factor) | |
| new_y_min = max(0, y_min - height * expand_factor) | |
| new_x_max = min(1, x_max + width * expand_factor) | |
| new_y_max = min(1, y_max + height * expand_factor) | |
| return (new_x_min, new_y_min, new_x_max, new_y_max) | |
| def inpaint_image(image_np, bbox, color=(255, 255, 255)): | |
| """ | |
| Inpaints a specified bounding box area in an image with a solid color. | |
| image_np: NumPy array of the image (BGR format). | |
| bbox: Tuple (x_min, y_min, x_max, y_max) in relative coordinates (0 to 1). | |
| color: Tuple (B, G, R) for the inpaint color. | |
| """ | |
| h, w, _ = image_np.shape | |
| x1 = int(bbox[0] * w) | |
| y1 = int(bbox[1] * h) | |
| x2 = int(bbox[2] * w) | |
| y2 = int(bbox[3] * h) | |
| # Create a copy to avoid modifying the original image array directly | |
| inpainted_image = image_np.copy() | |
| # Fill the bounding box with the specified color | |
| cv2.rectangle(inpainted_image, (x1, y1), (x2, y2), color, -1) # -1 fills the rectangle | |
| return inpainted_image | |
| # --- End Placeholder functions --- | |
| def process_pdf(pdf_file): | |
| """ | |
| Processes a PDF file, inpainting specified areas on each page. | |
| """ | |
| try: | |
| # Use test.pdf if nothing uploaded | |
| # Ensure 'wifi_basics_teaching_genspark.pdf' exists in the same directory | |
| # as your script for this to work when no file is uploaded. | |
| pdf_path = pdf_file.name if pdf_file else "wifi_basics_teaching_genspark.pdf" | |
| assert os.path.exists(pdf_path), f"File not found: {pdf_path}. Please upload a PDF or ensure 'wifi_basics_teaching_genspark.pdf' is in the current directory." | |
| # Define inpaint colors per page | |
| page_inpaint_colors = { | |
| 0: (255, 255, 255), # white for page 0 (BGR) | |
| 1: (0, 0, 0) # black for page 1 (BGR) | |
| } | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| # Convert PDF pages to images | |
| images = convert_from_path(pdf_path, dpi=300, output_folder=tmpdir, fmt="png") | |
| output_images = [] | |
| for i, img in enumerate(images): | |
| # Convert PIL RGB to OpenCV BGR | |
| np_img = np.array(img)[:, :, ::-1] | |
| # Expand the default bounding box for inpainting | |
| bbox = expand_bbox(default_bbox, expand_factor=0.3) | |
| # Get color for the current page, default to white if not specified | |
| color = page_inpaint_colors.get(i, (255, 255, 255)) | |
| # Perform inpainting | |
| inpainted = inpaint_image(np_img, bbox, color=color) | |
| # Save the inpainted image | |
| img_path = os.path.join(tmpdir, f"page_{i}.png") | |
| cv2.imwrite(img_path, inpainted) | |
| output_images.append(img_path) | |
| # Convert inpainted images back to a single PDF | |
| output_pdf = os.path.join(tmpdir, "inpainted_output.pdf") | |
| # Sort images by path to ensure correct page order in the output PDF | |
| with open(output_pdf, "wb") as f: | |
| f.write(img2pdf.convert([Path(p) for p in sorted(output_images)])) | |
| return output_pdf | |
| except Exception as e: | |
| print(f"❌ Error during PDF processing: {e}") | |
| # Return None to indicate an error, which safe_process_wrapper will handle | |
| return None | |
| # Gradio app UI | |
| with gr.Blocks() as demo: | |
| gr.Markdown("### 🧽 PDF Logo Inpainting (Solid Fill, Expand + Color Config)") | |
| with gr.Row(): | |
| # Input component for PDF upload | |
| pdf_input = gr.File( | |
| label="Upload PDF (or leave empty to use test.pdf)", | |
| file_types=[".pdf"], | |
| interactive=True | |
| ) | |
| # Output component for the processed PDF | |
| pdf_output = gr.File(label="Download Inpainted PDF") | |
| # Button to trigger the processing | |
| run_button = gr.Button("Process and Inpaint") | |
| # Link the button click to the processing function | |
| run_button.click( | |
| fn=process_pdf, | |
| inputs=pdf_input, | |
| outputs=pdf_output | |
| ) | |
| # Launch the Gradio application | |
| demo.launch() | |