import fitz # PyMuPDF import gradio as gr import zipfile import io import os from datetime import datetime import tempfile # Helper function to find color areas 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)) # Update the progress as each page is processed 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) } ] # Add custom color-comment pair if provided if custom_color and custom_comment and custom_stroke_color: custom_color_tuple = tuple(map(int, custom_color.split(','))) # Convert color to tuple custom_stroke_tuple = tuple(map(int, custom_stroke_color.split(','))) # Convert stroke to tuple color_comment_pairs.append({ "color": custom_color_tuple, "comment": custom_comment, "stroke_color": custom_stroke_tuple }) # Get the selected color-comment pairs selected_color_comment_pairs = [color_comment_pairs[i] for i in selected_color_comment_indices] # List to keep track of all modified PDFs (in-memory) modified_pdfs = [] # Process each input PDF file 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))) # Save the modified PDF in memory pdf_in_memory = io.BytesIO() doc.save(pdf_in_memory) doc.close() # Move pointer to the beginning of the BytesIO object pdf_in_memory.seek(0) modified_pdfs.append((os.path.basename(pdf_file.name), pdf_in_memory)) # Create an in-memory ZIP file containing all modified PDFs 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()) # Move pointer to the beginning of the ZIP file BytesIO object zip_in_memory.seek(0) # Generate the current date string current_date = datetime.now().strftime("%Y-%m-%d") # Create a named temporary file to store the ZIP 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 # Return the path to the temp file for Gradio to handle the download # Define the Gradio interface interface = gr.Interface( fn=process_pdf_files, inputs=[ gr.Files(label="Input PDF Files", file_types=[".pdf"]), gr.CheckboxGroup( # Use CheckboxGroup for multi-select 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"), # Use gr.File to trigger the download of the ZIP file title="PDF Color Region Markup" ) # Launch the Gradio app interface.launch()