qqwjq1981 commited on
Commit
9e369cd
·
verified ·
1 Parent(s): 114bbcb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -51
app.py CHANGED
@@ -4,15 +4,15 @@ import tempfile
4
  import numpy as np
5
  import cv2
6
  import img2pdf
 
 
7
  from pathlib import Path
8
- from pdf2image import convert_from_path # Make sure pdf2image is installed: pip install pdf2image opencv-python numpy img2pdf
9
 
10
- # Define a default bounding box (e.g., for a logo area)
11
  # Bounding box: (x1, y1, x2, y2) in fractions of image size
12
  default_bbox = (0.75, 0.9, 0.98, 0.98)
13
 
14
  def expand_bbox(bbox, expand_factor=0.3):
15
- """Expands a bounding box by a given factor."""
16
  x_min, y_min, x_max, y_max = bbox
17
  width = x_max - x_min
18
  height = y_max - y_min
@@ -24,99 +24,64 @@ def expand_bbox(bbox, expand_factor=0.3):
24
  return (new_x_min, new_y_min, new_x_max, new_y_max)
25
 
26
  def inpaint_image(image_np, bbox, color=(255, 255, 255)):
27
- """
28
- Inpaints a specified bounding box area in an image with a solid color.
29
- image_np: NumPy array of the image (BGR format).
30
- bbox: Tuple (x_min, y_min, x_max, y_max) in relative coordinates (0 to 1).
31
- color: Tuple (B, G, R) for the inpaint color.
32
- """
33
  h, w, _ = image_np.shape
34
  x1 = int(bbox[0] * w)
35
  y1 = int(bbox[1] * h)
36
  x2 = int(bbox[2] * w)
37
  y2 = int(bbox[3] * h)
38
-
39
- # Create a copy to avoid modifying the original image array directly
40
  inpainted_image = image_np.copy()
41
-
42
- # Fill the bounding box with the specified color
43
- cv2.rectangle(inpainted_image, (x1, y1), (x2, y2), color, -1) # -1 fills the rectangle
44
  return inpainted_image
45
 
46
- # --- End Placeholder functions ---
47
-
48
-
49
  def process_pdf(pdf_file):
50
- """
51
- Processes a PDF file, inpainting specified areas on each page.
52
- """
53
  try:
54
- # Use test.pdf if nothing uploaded
55
- # Ensure 'wifi_basics_teaching_genspark.pdf' exists in the same directory
56
- # as your script for this to work when no file is uploaded.
57
  pdf_path = pdf_file.name if pdf_file else "wifi_basics_teaching_genspark.pdf"
58
- 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."
59
 
60
- # Define inpaint colors per page
61
  page_inpaint_colors = {
62
- 0: (255, 255, 255), # white for page 0 (BGR)
63
- 1: (0, 0, 0) # black for page 1 (BGR)
64
  }
65
 
66
  with tempfile.TemporaryDirectory() as tmpdir:
67
- # Convert PDF pages to images
68
  images = convert_from_path(pdf_path, dpi=300, output_folder=tmpdir, fmt="png")
69
  output_images = []
70
 
71
  for i, img in enumerate(images):
72
- # Convert PIL RGB to OpenCV BGR
73
  np_img = np.array(img)[:, :, ::-1]
74
- # Expand the default bounding box for inpainting
75
  bbox = expand_bbox(default_bbox, expand_factor=0.3)
76
- # Get color for the current page, default to white if not specified
77
  color = page_inpaint_colors.get(i, (255, 255, 255))
78
- # Perform inpainting
79
  inpainted = inpaint_image(np_img, bbox, color=color)
80
 
81
- # Save the inpainted image
82
  img_path = os.path.join(tmpdir, f"page_{i}.png")
83
  cv2.imwrite(img_path, inpainted)
84
  output_images.append(img_path)
85
 
86
- # Convert inpainted images back to a single PDF
87
- output_pdf = os.path.join(tmpdir, "inpainted_output.pdf")
88
- # Sort images by path to ensure correct page order in the output PDF
89
- with open(output_pdf, "wb") as f:
90
  f.write(img2pdf.convert([Path(p) for p in sorted(output_images)]))
91
 
92
- return output_pdf
 
 
 
 
93
  except Exception as e:
94
  print(f"❌ Error during PDF processing: {e}")
95
- # Return None to indicate an error, which safe_process_wrapper will handle
96
  return None
97
 
98
- # Gradio app UI
99
  with gr.Blocks() as demo:
100
  gr.Markdown("### 🧽 PDF Logo Inpainting (Solid Fill, Expand + Color Config)")
101
  with gr.Row():
102
- # Input component for PDF upload
103
  pdf_input = gr.File(
104
  label="Upload PDF (or leave empty to use test.pdf)",
105
  file_types=[".pdf"],
106
  interactive=True
107
  )
108
- # Output component for the processed PDF
109
  pdf_output = gr.File(label="Download Inpainted PDF")
110
 
111
- # Button to trigger the processing
112
  run_button = gr.Button("Process and Inpaint")
 
113
 
114
- # Link the button click to the processing function
115
- run_button.click(
116
- fn=process_pdf,
117
- inputs=pdf_input,
118
- outputs=pdf_output
119
- )
120
-
121
- # Launch the Gradio application
122
  demo.launch()
 
4
  import numpy as np
5
  import cv2
6
  import img2pdf
7
+ import shutil
8
+ import uuid
9
  from pathlib import Path
10
+ from pdf2image import convert_from_path # pip install pdf2image opencv-python numpy img2pdf
11
 
 
12
  # Bounding box: (x1, y1, x2, y2) in fractions of image size
13
  default_bbox = (0.75, 0.9, 0.98, 0.98)
14
 
15
  def expand_bbox(bbox, expand_factor=0.3):
 
16
  x_min, y_min, x_max, y_max = bbox
17
  width = x_max - x_min
18
  height = y_max - y_min
 
24
  return (new_x_min, new_y_min, new_x_max, new_y_max)
25
 
26
  def inpaint_image(image_np, bbox, color=(255, 255, 255)):
 
 
 
 
 
 
27
  h, w, _ = image_np.shape
28
  x1 = int(bbox[0] * w)
29
  y1 = int(bbox[1] * h)
30
  x2 = int(bbox[2] * w)
31
  y2 = int(bbox[3] * h)
 
 
32
  inpainted_image = image_np.copy()
33
+ cv2.rectangle(inpainted_image, (x1, y1), (x2, y2), color, -1)
 
 
34
  return inpainted_image
35
 
 
 
 
36
  def process_pdf(pdf_file):
 
 
 
37
  try:
 
 
 
38
  pdf_path = pdf_file.name if pdf_file else "wifi_basics_teaching_genspark.pdf"
39
+ assert os.path.exists(pdf_path), f"File not found: {pdf_path}"
40
 
 
41
  page_inpaint_colors = {
42
+ 0: (255, 255, 255),
43
+ 1: (0, 0, 0)
44
  }
45
 
46
  with tempfile.TemporaryDirectory() as tmpdir:
 
47
  images = convert_from_path(pdf_path, dpi=300, output_folder=tmpdir, fmt="png")
48
  output_images = []
49
 
50
  for i, img in enumerate(images):
 
51
  np_img = np.array(img)[:, :, ::-1]
 
52
  bbox = expand_bbox(default_bbox, expand_factor=0.3)
 
53
  color = page_inpaint_colors.get(i, (255, 255, 255))
 
54
  inpainted = inpaint_image(np_img, bbox, color=color)
55
 
 
56
  img_path = os.path.join(tmpdir, f"page_{i}.png")
57
  cv2.imwrite(img_path, inpainted)
58
  output_images.append(img_path)
59
 
60
+ temp_pdf_path = os.path.join(tmpdir, "inpainted_output.pdf")
61
+ with open(temp_pdf_path, "wb") as f:
 
 
62
  f.write(img2pdf.convert([Path(p) for p in sorted(output_images)]))
63
 
64
+ # Save to persistent path
65
+ final_pdf_path = f"inpainted_output_{uuid.uuid4().hex[:8]}.pdf"
66
+ shutil.copy(temp_pdf_path, final_pdf_path)
67
+ return final_pdf_path
68
+
69
  except Exception as e:
70
  print(f"❌ Error during PDF processing: {e}")
 
71
  return None
72
 
73
+ # Gradio UI
74
  with gr.Blocks() as demo:
75
  gr.Markdown("### 🧽 PDF Logo Inpainting (Solid Fill, Expand + Color Config)")
76
  with gr.Row():
 
77
  pdf_input = gr.File(
78
  label="Upload PDF (or leave empty to use test.pdf)",
79
  file_types=[".pdf"],
80
  interactive=True
81
  )
 
82
  pdf_output = gr.File(label="Download Inpainted PDF")
83
 
 
84
  run_button = gr.Button("Process and Inpaint")
85
+ run_button.click(fn=process_pdf, inputs=pdf_input, outputs=pdf_output)
86
 
 
 
 
 
 
 
 
 
87
  demo.launch()