8.26 kB
| import os | |
| import pathlib | |
| import tempfile | |
| import time | |
| from io import BytesIO | |
| import pandas as pd | |
| import altair as alt | |
| import fitz | |
| import gradio as gr | |
| import PIL | |
| import skimage | |
| import streamlit as st | |
| from fastai.learner import load_learner | |
| from fastai.vision.all import * | |
| from fpdf import FPDF | |
| from icevision.all import * | |
| from icevision.models.checkpoint import * | |
| from PIL import Image as PILImage | |
| CHECKPOINT_PATH = "./allsynthetic-imgsize768.pth" | |
| def load_icevision_model(): | |
| return model_from_checkpoint(CHECKPOINT_PATH) | |
| def load_fastai_model(): | |
| return load_learner("fastai-classification-model.pkl") | |
| checkpoint_and_model = load_icevision_model() | |
| model = checkpoint_and_model["model"] | |
| model_type = checkpoint_and_model["model_type"] | |
| class_map = checkpoint_and_model["class_map"] | |
| img_size = checkpoint_and_model["img_size"] | |
| valid_tfms = tfms.A.Adapter( | |
| [*tfms.A.resize_and_pad(img_size), tfms.A.Normalize()] | |
| ) | |
| learn = load_fastai_model() | |
| labels = learn.dls.vocab | |
| def get_content_area(pred_dict) -> int: | |
| if "content" not in pred_dict["detection"]["labels"]: | |
| return 0 | |
| content_bboxes = [ | |
| pred_dict["detection"]["bboxes"][idx] | |
| for idx, label in enumerate(pred_dict["detection"]["labels"]) | |
| if label == "content" | |
| ] | |
| cb = content_bboxes[0] | |
| return (cb.xmax - cb.xmin) * (cb.ymax - cb.ymin) | |
| def get_redaction_area(pred_dict) -> int: | |
| if "redaction" not in pred_dict["detection"]["labels"]: | |
| return 0 | |
| redaction_bboxes = [ | |
| pred_dict["detection"]["bboxes"][idx] | |
| for idx, label in enumerate(pred_dict["detection"]["labels"]) | |
| if label == "redaction" | |
| ] | |
| return sum( | |
| (bbox.xmax - bbox.xmin) * (bbox.ymax - bbox.ymin) | |
| for bbox in redaction_bboxes | |
| ) | |
| st.title("Redaction Detector") | |
| st.image( | |
| "./synthetic-redactions.jpg", | |
| width=300, | |
| ) | |
| uploaded_pdf = st.file_uploader( | |
| "Upload a PDF...", | |
| type="pdf", | |
| accept_multiple_files=False, | |
| help="This application processes PDF files. Please upload a document you believe to contain redactions.", | |
| on_change=None, | |
| ) | |
| # Add a selectbox to the sidebar: | |
| st.sidebar.header("Customisation Options") | |
| graph_checkbox = st.sidebar.checkbox( | |
| "Show analysis charts", | |
| value=True, | |
| help="Display charts analysising the redactions found in the document.", | |
| ) | |
| extract_images_checkbox = st.sidebar.checkbox( | |
| "Extract redacted images", | |
| value=True, | |
| help="Create a PDF file containing the redacted images with an object detection overlay highlighting their locations and the confidence the model had when detecting the redactions.", | |
| ) | |
| # Add a slider to the sidebar: | |
| confidence = st.sidebar.slider( | |
| "Confidence level (%)", | |
| min_value=0, | |
| max_value=100, | |
| value=80, | |
| ) | |
| def get_pdf_document(input): | |
| os.makedirs(str(pathlib.Path(filename_without_extension)), exist_ok=True) | |
| with open( | |
| str(pathlib.Path(filename_without_extension) / "output.pdf"), "wb" | |
| ) as f: | |
| f.write(uploaded_pdf.getbuffer()) | |
| return fitz.open( | |
| str(pathlib.Path(filename_without_extension) / "output.pdf") | |
| ) | |
| def get_image_predictions(img): | |
| return model_type.end2end_detect( | |
| img, | |
| valid_tfms, | |
| model, | |
| class_map=class_map, | |
| detection_threshold=confidence / 100, | |
| display_label=True, | |
| display_bbox=True, | |
| return_img=True, | |
| font_size=16, | |
| label_color="#FF59D6", | |
| ) | |
| if uploaded_pdf is None: | |
| st.markdown(pathlib.Path("article.md").read_text()) | |
| else: | |
| st.text("Opening PDF...") | |
| filename_without_extension = uploaded_pdf.name[:-4] | |
| results = [] | |
| images = [] | |
| document = get_pdf_document(uploaded_pdf) | |
| total_image_areas = 0 | |
| total_content_areas = 0 | |
| total_redaction_area = 0 | |
| tmp_dir = tempfile.gettempdir() | |
| for page_num, page in enumerate(document, start=1): | |
| image_pixmap = page.get_pixmap() | |
| image = image_pixmap.tobytes() | |
| _, _, probs = learn.predict(image) | |
| results.append( | |
| {labels[i]: float(probs[i]) for i in range(len(labels))} | |
| ) | |
| if probs[0] > (confidence / 100): | |
| redaction_count = len(images) | |
| if not os.path.exists( | |
| os.path.join(tmp_dir, filename_without_extension or "abc") | |
| ): | |
| os.makedirs(os.path.join(tmp_dir, filename_without_extension)) | |
| image_pixmap.save( | |
| os.path.join( | |
| tmp_dir, filename_without_extension, f"page-{page_num}.png" | |
| ) | |
| ) | |
| images.append( | |
| [ | |
| f"Redacted page #{redaction_count + 1} on page {page_num}", | |
| os.path.join( | |
| tmp_dir, | |
| filename_without_extension, | |
| f"page-{page_num}.png", | |
| ), | |
| ] | |
| ) | |
| redacted_pages = [ | |
| str(page + 1) | |
| for page in range(len(results)) | |
| if results[page]["redacted"] > (confidence / 100) | |
| ] | |
| report = os.path.join( | |
| tmp_dir, filename_without_extension, "redacted_pages.pdf" | |
| ) | |
| if extract_images_checkbox: | |
| with st.spinner('Calculating redaction proportions...'): | |
| pdf = FPDF(unit="cm", format="A4") | |
| pdf.set_auto_page_break(0) | |
| imagelist = sorted( | |
| [ | |
| i | |
| for i in os.listdir( | |
| os.path.join(tmp_dir, filename_without_extension) | |
| ) | |
| if i.endswith("png") | |
| ] | |
| ) | |
| for image in imagelist: | |
| with PILImage.open( | |
| os.path.join(tmp_dir, filename_without_extension, image) | |
| ) as img: | |
| size = img.size | |
| width, height = size | |
| if width > height: | |
| pdf.add_page(orientation="L") | |
| else: | |
| pdf.add_page(orientation="P") | |
| pred_dict = get_image_predictions(img) | |
| total_image_areas += pred_dict["width"] * pred_dict["height"] | |
| total_content_areas += get_content_area(pred_dict) | |
| total_redaction_area += get_redaction_area(pred_dict) | |
| pred_dict["img"].save( | |
| os.path.join( | |
| tmp_dir, filename_without_extension, f"pred-{image}" | |
| ), | |
| ) | |
| pdf.image( | |
| os.path.join( | |
| tmp_dir, filename_without_extension, f"pred-{image}" | |
| ), | |
| w=pdf.w, | |
| h=pdf.h, | |
| ) | |
| pdf.output(report, "F") | |
| st.success('Image predictions complete!') | |
| text_output = f"A total of {len(redacted_pages)} pages were redacted. \n\nThe redacted page numbers were: {', '.join(redacted_pages)}. \n\n" | |
| st.balloons() | |
| if not extract_images_checkbox: | |
| st.text(text_output) | |
| # DISPLAY IMAGES | |
| else: | |
| total_redaction_proportion = round( | |
| (total_redaction_area / total_image_areas) * 100, 1 | |
| ) | |
| content_redaction_proportion = round( | |
| (total_redaction_area / total_content_areas) * 100, 1 | |
| ) | |
| redaction_analysis = f"- {total_redaction_proportion}% of the total area of the redacted pages was redacted. \n- {content_redaction_proportion}% of the actual content of those redacted pages was redacted." | |
| source = pd.DataFrame( | |
| { | |
| "category": ["Unredacted", "Redacted"], | |
| "value": [ | |
| 100 - total_redaction_proportion, | |
| total_redaction_proportion, | |
| ], | |
| } | |
| ) | |
| c = ( | |
| alt.Chart(source) | |
| .mark_arc() | |
| .encode( | |
| theta=alt.Theta(field="value", type="quantitative"), | |
| color=alt.Color(field="category", type="nominal"), | |
| ) | |
| ) | |
| st.altair_chart(c, use_container_width=True) | |
| st.text(text_output + redaction_analysis) | |
| # DISPLAY IMAGES | |