File size: 5,445 Bytes
e982db7 b922398 e982db7 b922398 e982db7 b922398 e982db7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | import streamlit as st
import pandas as pd
from PyPDF2 import PdfReader
from difflib import ndiff
import os
import tempfile
import matplotlib.pyplot as plt
from doctr.io import DocumentFile
from doctr.models import ocr_predictor
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import fitz
# Function to find similar paragraphs using cosine similarity
def find_similar(vector_representation, all_representations, threshold=0.5):
similarity_matrix = cosine_similarity(vector_representation, all_representations)
similarities = similarity_matrix[0]
# Find indices where similarity is above the threshold
similar_indices = np.where(similarities > threshold)[0]
return similar_indices.tolist()
def process_pdf(doc, diff_line):
data_geom = []
data_word = []
data_page = []
for page_num, page in enumerate(doc.pages()):
for block in page.get_text("blocks"):
for line in block["lines"]:
temp = []
data_page.append(page_num)
data_geom.append(line["bbox"])
for word in line["words"]:
temp.append(word["text"])
data_word.append(" ".join(temp))
data = {
"halaman": data_page,
"kalimat": data_word,
"posisi": data_geom
}
df = pd.DataFrame(data)
y1_coor = []
y2_coor = []
for pos in df["posisi"]:
y1_coor.append(pos[1])
y2_coor.append(pos[3])
df["y1_coor"] = y1_coor
df["y2_coor"] = y2_coor
df = df.sort_values(by=["halaman", "y2_coor"])
temp_kalimat = []
merge_kalimat = []
merge_y1 = []
merge_y2 = []
merge_halaman = []
for i in range(df.shape[0]):
if i == 0:
temp_y2 = df["y2_coor"].iloc[i]
temp_y1 = df["y1_coor"].iloc[i]
temp_hal = df["halaman"].iloc[i]
temp_kalimat.append(df["kalimat"].iloc[i])
else:
if abs(temp_y2 - df["y1_coor"].iloc[i]) < diff_line:
temp_kalimat.append(df["kalimat"].iloc[i])
temp_y2 = df["y2_coor"].iloc[i]
else:
merge_kalimat.append(" ".join(temp_kalimat))
merge_y1.append(temp_y1)
merge_y2.append(temp_y2)
merge_halaman.append(temp_hal)
temp_y2 = df["y2_coor"].iloc[i]
temp_y1 = df["y1_coor"].iloc[i]
temp_hal = df["halaman"].iloc[i]
temp_kalimat = []
temp_kalimat.append(df["kalimat"].iloc[i])
data = {
"halaman": merge_halaman,
"kalimat": merge_kalimat,
"y1_coor": merge_y1,
"y2_coor": merge_y2
}
return pd.DataFrame(data)
def highlight_pdf(doc, delete_df):
for i in range(delete_df.shape[0]):
norm_y1 = delete_df["y1_coor"].iloc[i]
norm_y2 = delete_df["y2_coor"].iloc[i]
page = int(delete_df["halaman"].iloc[i])
doc[page].draw_rect(
fitz.Rect(1, norm_y1, doc[page].rect.width, norm_y2),
color=(0, 1, 0),
width=2
)
def get_download_link(file_path, text):
"""Generate a download link for the given file."""
href = f'<a href="{file_path}" download="{file_path}">{text}</a>'
return href
# Main Streamlit app
def main():
st.title("PDF Paragraph Highlighter")
# Upload PDF file
uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])
# Sidebar for threshold and line_diff values
threshold = st.sidebar.number_input("Threshold", min_value=0.1, max_value=1.0, value=0.4, step=0.1)
line_diff = st.sidebar.number_input("Line Diff Value", min_value=1.0, max_value=100.0, value=50.0, step=5.0)
line_diff = line_diff / 10000
if uploaded_file is not None:
st.subheader("Uploaded PDF Preview")
st.write(f"Filename: {uploaded_file.name}")
# Processing PDF content
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(uploaded_file.read())
temp_file_path = temp_file.name
doc = fitz.open(temp_file_path)
res_df = process_pdf(doc, line_diff)
# Display processed DataFrame
st.subheader("Processed DataFrame")
st.dataframe(res_df)
# Search for similar paragraphs
search_string = "deletion erasure removal purge expunge wipe clear annulment obliteration destruction"
embeddings_distilbert = model.encode(res_df['kalimat'].values)
search_vect = model.encode([search_string])
distilbert_similar_indexes = find_similar(search_vect, embeddings_distilbert, threshold)
delete_df = res_df.iloc[distilbert_similar_indexes, :]
# Highlight paragraphs in the PDF
highlight_pdf(doc, delete_df)
# Save the highlighted PDF
highlighted_path = './highlighted.pdf'
doc.save(highlighted_path)
# Display download link for the highlighted PDF
st.subheader("Download Highlighted PDF")
st.markdown(get_download_link(highlighted_path, 'Download Highlighted PDF'), unsafe_allow_html=True)
# Cleanup: Remove temporary file
os.remove(temp_file_path)
if __name__ == "__main__":
# Load the SentenceTransformer model
model = SentenceTransformer('distilbert-base-nli-stsb-mean-tokens')
# Run the Streamlit app
main()
|