import os import requests from PIL import Image, ImageDraw import io import base64 import json import gradio as gr import fitz # PyMuPDF import tempfile from typing import Union # --- Configuration & API Constants --- INVOKE_URL_OCR = "https://ai.api.nvidia.com/v1/cv/nvidia/nemoretriever-ocr-v1" INVOKE_URL_PARSER = "https://integrate.api.nvidia.com/v1/chat/completions" MAX_PIXELS_FOR_PARSER = 1024 * 1024 # 1 Megapixel # ================================================================================= # SELF-CONTAINED REDACTION LOGIC # (This is the refined function from the previous step) # ================================================================================= def _get_average_color_from_regions(image: Image.Image, regions: list[tuple]): """Calculates the average RGB color from a list of regions in an image.""" total_r, total_g, total_b = 0, 0, 0; pixel_count = 0 img_width, img_height = image.size if image.mode == 'RGBA': image = image.convert('RGB') pixels = image.load() for region in regions: x1, y1, x2, y2 = [max(0, int(c)) for c in region] x2 = min(img_width, x2); y2 = min(img_height, y2) for x in range(x1, x2): for y in range(y1, y2): r, g, b = pixels[x, y] total_r += r; total_g += g; total_b += b pixel_count += 1 if pixel_count == 0: return (0, 0, 0) return (total_r // pixel_count, total_g // pixel_count, total_b // pixel_count) def _detect_pictures_with_parser(image_to_process: Image.Image, api_key: str): """Sends an image to the NemoRetriever Parser model to detect 'Picture' elements.""" headers = {"Authorization": f"Bearer {api_key}", "Accept": "application/json"} buffered = io.BytesIO() image_to_process.save(buffered, format="PNG") b64_str = base64.b64encode(buffered.getvalue()).decode("ascii") content = f'' tool_name = "markdown_bbox" payload = { "model": "nvidia/nemoretriever-parse", "messages": [{"role": "user", "content": content}], "tools": [{"type": "function", "function": {"name": tool_name}}], "tool_choice": {"type": "function", "function": {"name": tool_name}}, "max_tokens": 2048, } response = requests.post(INVOKE_URL_PARSER, headers=headers, json=payload, timeout=120) response.raise_for_status() response_json = response.json() picture_bboxes = [] tool_calls = response_json.get('choices', [{}])[0].get('message', {}).get('tool_calls', []) if tool_calls: arguments_str = tool_calls[0].get('function', {}).get('arguments', '[]') parsed_arguments = json.loads(arguments_str) if parsed_arguments and isinstance(parsed_arguments, list): for element in parsed_arguments[0]: if element.get("type") == "Picture" and element.get("bbox"): picture_bboxes.append(element["bbox"]) return picture_bboxes def _redact_text_in_image(input_image: Image.Image, api_key: str): """Sends a (cropped) image to the OCR model and returns a redacted version.""" headers = {"Authorization": f"Bearer {api_key}", "Accept": "application/json"} buffered = io.BytesIO(); input_image.save(buffered, format="PNG") image_b64 = base64.b64encode(buffered.getvalue()).decode() payload = {"input": [{"type": "image_url", "url": f"data:image/png;base64,{image_b64}"}]} try: response = requests.post(INVOKE_URL_OCR, headers=headers, json=payload, timeout=60) response.raise_for_status(); response_json = response.json() except requests.exceptions.RequestException: return input_image image_with_redactions = input_image.copy(); draw = ImageDraw.Draw(image_with_redactions) img_width, img_height = image_with_redactions.size radius = max(1, int(((img_width**2 + img_height**2)**0.5) / 100)) try: detections = response_json['data'][0]['text_detections'] for detection in detections: bbox = detection.get("bounding_box") if bbox and bbox.get("points"): points = bbox["points"] p1 = (points[0]['x'] * img_width, points[0]['y'] * img_height); p3 = (points[2]['x'] * img_width, points[2]['y'] * img_height) sample_regions = [(p1[0], p1[1] - radius, p3[0], p1[1]), (p1[0], p3[1], p3[0], p3[1] + radius), (p1[0] - radius, p1[1], p1[0], p3[1]), (p3[0], p1[1], p3[0] + radius, p3[1])] redaction_color = _get_average_color_from_regions(image_with_redactions, sample_regions) draw.rectangle([p1, p3], fill=redaction_color) return image_with_redactions except (KeyError, IndexError, TypeError): return input_image def redact_pictures_in_image(image_source: Union[str, Image.Image], api_key: str, callback: callable = None) -> Image.Image: """ Analyzes an image to find pictures, then redacts text within those pictures. Now accepts a file path, base64 string, or a PIL Image object directly. """ def _progress(message: str): if callback: callback(message) _progress("Loading image for processing...") try: if isinstance(image_source, Image.Image): input_image = image_source.convert("RGB") elif os.path.exists(image_source): input_image = Image.open(image_source).convert("RGB") else: input_image = Image.open(io.BytesIO(base64.b64decode(image_source))).convert("RGB") except Exception as e: raise ValueError(f"Invalid image_source. Error: {e}") image_to_analyze = input_image original_width, original_height = input_image.size if (original_width * original_height) > MAX_PIXELS_FOR_PARSER: _progress(f"Image is large, resizing for analysis...") scale = (MAX_PIXELS_FOR_PARSER / (original_width * original_height))**0.5 new_dims = (int(original_width * scale), int(original_height * scale)) image_to_analyze = input_image.resize(new_dims, Image.Resampling.LANCZOS) _progress("Detecting 'Picture' elements...") try: picture_bboxes = _detect_pictures_with_parser(image_to_analyze, api_key) except requests.exceptions.RequestException as e: _progress(f"API Error during picture detection: {e}"); raise if not picture_bboxes: _progress("No 'Picture' elements found.") return input_image _progress(f"Found {len(picture_bboxes)} 'Picture' element(s). Redacting text...") final_image = input_image.copy() for i, box in enumerate(picture_bboxes): _progress(f" - Processing picture {i + 1}/{len(picture_bboxes)}...") x1, y1 = int(box["xmin"] * original_width), int(box["ymin"] * original_height) x2, y2 = int(box["xmax"] * original_width), int(box["ymax"] * original_height) cropped_element = input_image.crop((x1, y1, x2, y2)) redacted_crop = _redact_text_in_image(cropped_element, api_key) final_image.paste(redacted_crop, (x1, y1)) _progress("Redaction for this page complete.") return final_image # ================================================================================= # GRADIO PDF PROCESSING APPLICATION # ================================================================================= def process_pdf(pdf_file, progress=gr.Progress(track_tqdm=True)): """ Main function for the Gradio app. Takes an uploaded PDF file, processes each page, and returns the path to the redacted output PDF. """ if pdf_file is None: raise gr.Error("Please upload a PDF file.") api_key = os.getenv("NVIDIA_API_KEY") if not api_key: raise gr.Error("NVIDIA_API_KEY environment variable not set.") log_messages = [] def progress_callback(message): print(message) # Also print to console for debugging log_messages.append(message) try: pdf_path = pdf_file.name doc = fitz.open(pdf_path) processed_pages = [] for page_num in progress.tqdm(range(len(doc)), desc="Processing PDF Pages"): progress_callback(f"\n--- Processing Page {page_num + 1} of {len(doc)} ---") # Convert page to image (150 DPI is a good balance of quality and size) page = doc.load_page(page_num) pix = page.get_pixmap(dpi=150) page_image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) # Run the redaction pipeline on the single page image processed_image = redact_pictures_in_image( image_source=page_image, api_key=api_key, callback=progress_callback ) processed_pages.append(processed_image) progress_callback("\n--- Finalizing PDF ---") if not processed_pages: raise gr.Error("No pages were processed from the PDF.") # Save processed images into a new PDF output_pdf_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name processed_pages[0].save( output_pdf_path, "PDF", resolution=100.0, save_all=True, append_images=processed_pages[1:] ) progress_callback(f"Successfully created redacted PDF: {os.path.basename(output_pdf_path)}") return output_pdf_path, "\n".join(log_messages) except Exception as e: gr.Error(f"An error occurred: {e}") return None, f"An error occurred: {e}" # --- Gradio UI Definition --- if __name__ == "__main__": with gr.Blocks(theme=gr.themes.Default(), title="NVIDIA PDF Redactor") as demo: gr.Markdown( """ # document Redactor for Pictures Upload a PDF document. The tool will scan each page for pictures, redact any text found exclusively within those pictures, and then generate a new, downloadable PDF with the redactions. Pages without pictures are skipped to save time and cost. """ ) with gr.Row(): with gr.Column(scale=1): pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"]) process_btn = gr.Button("🚀 Process PDF and Redact Pictures", variant="primary") with gr.Column(scale=2): pdf_output = gr.File(label="Download Redacted PDF", interactive=False) status_log = gr.Textbox(label="Processing Log", lines=15, interactive=False) process_btn.click( fn=process_pdf, inputs=[pdf_input], outputs=[pdf_output, status_log] ) gr.Markdown("---") gr.Markdown("Powered by [NVIDIA NIM](https://build.nvidia.com/explore/discover).") demo.launch(debug=True)