| import fitz
|
| import gradio as gr
|
| import zipfile
|
| import io
|
| import os
|
| from datetime import datetime
|
| import tempfile
|
|
|
|
|
| def find_color_areas(page, target_color, tolerance=30):
|
| pix = page.get_pixmap()
|
| width, height = pix.width, pix.height
|
| visited = [[False for _ in range(width)] for _ in range(height)]
|
| rectangles = []
|
|
|
| def flood_fill(x, y):
|
| stack = [(x, y)]
|
| rects = []
|
| while stack:
|
| cx, cy = stack.pop()
|
| if visited[cy][cx]:
|
| continue
|
|
|
| visited[cy][cx] = True
|
| pixel_color = pix.pixel(cx, cy)
|
| r, g, b = pixel_color[:3]
|
|
|
| if (abs(r - target_color[0]) <= tolerance and
|
| abs(g - target_color[1]) <= tolerance and
|
| abs(b - target_color[2]) <= tolerance):
|
| rects.append(fitz.Rect(cx, cy, cx + 1, cy + 1))
|
| if cx > 0: stack.append((cx - 1, cy))
|
| if cx < width - 1: stack.append((cx + 1, cy))
|
| if cy > 0: stack.append((cx, cy - 1))
|
| if cy < height - 1: stack.append((cx, cy + 1))
|
|
|
| if rects:
|
| bbox = fitz.Rect(min([r.x0 for r in rects]),
|
| min([r.y0 for r in rects]),
|
| max([r.x1 for r in rects]),
|
| max([r.y1 for r in rects]))
|
| return bbox
|
| return None
|
|
|
| for y in range(height):
|
| for x in range(width):
|
| if not visited[y][x]:
|
| bbox = flood_fill(x, y)
|
| if bbox:
|
| rectangles.append(bbox)
|
| return rectangles
|
|
|
| def merge_overlapping_rectangles(rectangles):
|
| merged_rects = []
|
| while rectangles:
|
| rect = rectangles.pop(0)
|
| to_merge = [rect]
|
| for other in rectangles[:]:
|
| if rect.intersects(other):
|
| to_merge.append(other)
|
| rectangles.remove(other)
|
| merged_rect = fitz.Rect(
|
| min([r.x0 for r in to_merge]),
|
| min([r.y0 for r in to_merge]),
|
| max([r.x1 for r in to_merge]),
|
| max([r.y1 for r in to_merge])
|
| )
|
| merged_rects.append(merged_rect)
|
| return merged_rects
|
|
|
| def markup_color_regions(doc, color_comment_pairs, tolerance=30, progress=None):
|
| for page_num in range(len(doc)):
|
| page = doc[page_num]
|
| progress((page_num + 1) / len(doc))
|
| for color_comment_pair in color_comment_pairs:
|
| target_color = color_comment_pair['color']
|
| comment = color_comment_pair['comment']
|
| stroke_color = color_comment_pair['stroke_color']
|
|
|
| rectangles = find_color_areas(page, target_color, tolerance)
|
| if rectangles:
|
| merged_rectangles = merge_overlapping_rectangles(rectangles)
|
| for bbox in merged_rectangles:
|
| annot = page.add_rect_annot(bbox)
|
| annot.set_colors(stroke=stroke_color)
|
| annot.set_border(width=2)
|
| annot.set_info({"title": "Markup", "content": comment})
|
| annot.update()
|
|
|
| def process_pdf_files(input_pdfs, selected_color_comment_indices, tolerance, custom_color, custom_comment, custom_stroke_color, progress=gr.Progress()):
|
| color_comment_pairs = [
|
| {
|
| "color": (235, 128, 138),
|
| "comment": "Structural Slab greater than architectural slab",
|
| "stroke_color": (1, 0, 0)
|
| },
|
| {
|
| "color": (128, 253, 128),
|
| "comment": "Arch Slab greater than Structure",
|
| "stroke_color": (0, 1, 0)
|
| }
|
| ]
|
|
|
|
|
| if custom_color and custom_comment and custom_stroke_color:
|
| custom_color_tuple = tuple(map(int, custom_color.split(',')))
|
| custom_stroke_tuple = tuple(map(int, custom_stroke_color.split(',')))
|
| color_comment_pairs.append({
|
| "color": custom_color_tuple,
|
| "comment": custom_comment,
|
| "stroke_color": custom_stroke_tuple
|
| })
|
|
|
|
|
| selected_color_comment_pairs = [color_comment_pairs[i] for i in selected_color_comment_indices]
|
|
|
|
|
| modified_pdfs = []
|
|
|
|
|
| for i, pdf_file in enumerate(input_pdfs):
|
| with open(pdf_file.name, "rb") as file_stream:
|
| doc = fitz.open(stream=file_stream.read(), filetype="pdf")
|
| markup_color_regions(doc, selected_color_comment_pairs, tolerance, progress=lambda p: progress((i + p) / len(input_pdfs)))
|
|
|
|
|
| pdf_in_memory = io.BytesIO()
|
| doc.save(pdf_in_memory)
|
| doc.close()
|
|
|
|
|
| pdf_in_memory.seek(0)
|
| modified_pdfs.append((os.path.basename(pdf_file.name), pdf_in_memory))
|
|
|
|
|
| zip_in_memory = io.BytesIO()
|
| with zipfile.ZipFile(zip_in_memory, 'w') as zipf:
|
| for pdf_name, pdf_bytes in modified_pdfs:
|
| zipf.writestr(pdf_name, pdf_bytes.read())
|
|
|
|
|
| zip_in_memory.seek(0)
|
|
|
|
|
| current_date = datetime.now().strftime("%Y-%m-%d")
|
|
|
|
|
| temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=f"_CoordinationPDFS_{current_date}.zip")
|
| with open(temp_file.name, 'wb') as tmp_file:
|
| tmp_file.write(zip_in_memory.read())
|
|
|
| return temp_file.name
|
|
|
|
|
| interface = gr.Interface(
|
| fn=process_pdf_files,
|
| inputs=[
|
| gr.Files(label="Input PDF Files", file_types=[".pdf"]),
|
| gr.CheckboxGroup(
|
| label="Select Color-Comment Pairs",
|
| choices=["Structural Slab vs Arch Slab", "Arch Slab vs Structural Slab", "Custom Option"],
|
| type="index"
|
| ),
|
| gr.Slider(label="Tolerance", minimum=0, maximum=100, step=1, value=30),
|
| gr.Textbox(label="Custom Color (R,G,B)", placeholder="Enter custom color in RGB format, e.g., 255,0,0"),
|
| gr.Textbox(label="Custom Comment", placeholder="Enter custom comment for this color"),
|
| gr.Textbox(label="Custom Stroke Color (R,G,B)", placeholder="Enter stroke color in RGB format, e.g., 0,0,255")
|
| ],
|
| outputs=gr.File(label="Download ZIP File"),
|
| title="PDF Color Region Markup"
|
| )
|
|
|
|
|
| interface.launch()
|
|
|
