Update app.py

app.py

@@ -4,15 +4,15 @@ import tempfile
 import numpy as np
 import cv2
 import img2pdf
+import shutil
+import uuid
 from pathlib import Path
-from pdf2image import convert_from_path # Make sure pdf2image is installed: pip install pdf2image opencv-python numpy img2pdf
+from pdf2image import convert_from_path  # 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
@@ -24,99 +24,64 @@ def expand_bbox(bbox, expand_factor=0.3):
     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
+    cv2.rectangle(inpainted_image, (x1, y1), (x2, y2), color, -1)
     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."
+        assert os.path.exists(pdf_path), f"File not found: {pdf_path}"
 
-        # 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)
+            0: (255, 255, 255),
+            1: (0, 0, 0)
         }
 
         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:
+            temp_pdf_path = os.path.join(tmpdir, "inpainted_output.pdf")
+            with open(temp_pdf_path, "wb") as f:
                 f.write(img2pdf.convert([Path(p) for p in sorted(output_images)]))
 
-            return output_pdf
+        # Save to persistent path
+        final_pdf_path = f"inpainted_output_{uuid.uuid4().hex[:8]}.pdf"
+        shutil.copy(temp_pdf_path, final_pdf_path)
+        return final_pdf_path
+
     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
+# Gradio 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")
+    run_button.click(fn=process_pdf, inputs=pdf_input, outputs=pdf_output)
 
-    # 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()