Switch to Full Page Rasterization (Nuclear Method)

app.py

@@ -69,76 +69,68 @@ import fitz  # PyMuPDF
 from concurrent.futures import ThreadPoolExecutor
 
 def process_pdf(pdf_file, image_editor_data):
-    """Surgeon Method: Object Removal (No AI Patching)"""
+    """Nuclear Method: Full Page Rasterization"""
     if pdf_file is None or image_editor_data is None:
         return None
 
-    # 1. Prepare Mask & Coordinates
+    # 1. Get the mask defined by user on Page 1
     full_mask = get_mask_from_dict(image_editor_data)
     
     # Dilate mask slightly to be safe
     kernel = np.ones((5,5), np.uint8)
-    full_mask = cv2.dilate(full_mask, kernel, iterations=2)
-
-    # Find bounding box of the mask
-    rows, cols = np.where(full_mask > 0)
-    if len(rows) == 0:
-        return pdf_file.name 
-    
-    y1, y2 = np.min(rows), np.max(rows)
-    x1, x2 = np.min(cols), np.max(cols)
-    
-    # Open PDF
-    doc = fitz.open(pdf_file.name)
-    total_pages = len(doc)
+    full_mask = cv2.dilate(full_mask, kernel, iterations=3)
+
+    # 2. Convert ALL pages to High-Res Images (300 DPI)
+    # This "flattens" vector graphics into pixels, solving the color profile mismatch.
+    print("Rasterizing PDF to Images (300 DPI)...")
+    try:
+        pages = pdf2image.convert_from_path(pdf_file.name, dpi=300)
+    except Exception as e:
+        print(f"Error converting PDF: {e}")
+        return None
     
-    print(f"Surgeon Processing {total_pages} pages...")
+    cleaned_pages = []
+    total_pages = len(pages)
+    print(f"Processing {total_pages} pages...")
 
-    for i in range(total_pages):
-        page = doc[i]
+    for i, page in enumerate(pages):
+        # Convert PIL to Numpy
+        img_np = np.array(page)
         
-        # Map pixels (300 dpi) -> PDF Points (72 dpi)
-        scale = 72 / 300
+        # Ensure RGB
+        if len(img_np.shape) == 2:
+            img_np = cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB)
+        elif len(img_np.shape) == 3 and img_np.shape[2] == 4:
+            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2RGB)
+            
+        # Resize mask if page size differs from preview
+        if img_np.shape[:2] != full_mask.shape[:2]:
+            current_mask = cv2.resize(full_mask, (img_np.shape[1], img_np.shape[0]), interpolation=cv2.INTER_NEAREST)
+        else:
+            current_mask = full_mask
+
+        # Run AI (Inpainting)
+        # Since input is now RGB pixels, the AI's RGB output will blend much better.
+        result = lama.predict(img_np, current_mask)
         
-        # Define the removal area in PDF coordinates
-        rect = fitz.Rect(
-            x1 * scale,
-            y1 * scale,
-            x2 * scale,
-            y2 * scale
-        )
-        
-        # 1. Image Object Removal (For Logos)
-        # Scan for images that are fully inside our mask
-        images = page.get_images()
-        for img in images:
-            xref = img[0]
-            # Get image bbox
-            try:
-                img_rect = page.get_image_rect(xref)
-                # If image is mostly inside our delete-zone
-                if rect.intersects(img_rect):
-                    intersection = rect & img_rect
-                    if intersection.get_area() > (img_rect.get_area() * 0.8):
-                        print(f"Deleting logo/image on page {i+1}")
-                        page.delete_image(xref)
-            except:
-                pass
-
-        # 2. Text/Vector Removal (For Text Watermarks)
-        # We redact the area, but we tell it NOT to remove background images
-        page.add_redact_annot(rect, fill=[], stroke=[]) # Transparent redaction
-        
-        # CRITICAL FLAG: fitz.PDF_REDACT_IMAGE_NONE
-        # This tells it: "Remove text and drawings, but LEAVE THE BACKGROUND IMAGE ALONE"
-        page.apply_redactions(images=fitz.PDF_REDACT_IMAGE_NONE)
+        # Convert back to PIL for PDF saving
+        cleaned_pages.append(Image.fromarray(result))
+        print(f"Processed page {i+1}/{total_pages}")
 
-    # 5. Save
+    # 3. Save back to PDF with Max Quality
     output_path = tempfile.mktemp(suffix=".pdf")
-    doc.save(output_path, garbage=4, deflate=True)
-    doc.close()
+    if cleaned_pages:
+        cleaned_pages[0].save(
+            output_path, 
+            save_all=True, 
+            append_images=cleaned_pages[1:],
+            quality=100,      # Max JPEG quality
+            resolution=300.0, # Maintain High DPI
+            subsampling=0     # Disable chroma subsampling for sharper colors
+        )
+        return output_path
     
-    return output_path
+    return None
 
 
 # --- UI Construction ---