sankhyan/st_table_extractor
0
1# File: streamlit_app.py
2
3import os
4import re
5import csv
6import tempfile
7import base64
8
9import pandas as pd
10import PyPDF2
11import pytesseract
12from pdf2image import convert_from_path
13import streamlit as st
14
15# ─── Configuration ────────────────────────────────────────────────────────────
16
17POPPLER_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "poppler", "bin")
18TESSERACT_CMD = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
19pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD
20
21
22# ─── Helper Functions (from your original Flask code) ─────────────────────────
23
24def get_poppler_path():
25 return POPPLER_PATH
26
27def convert_hindi_digits(text: str) -> str:
28 hindi_digits = {
29 "०": "0", "१": "1", "२": "2", "३": "3", "४": "4",
30 "५": "5", "६": "6", "७": "7", "८": "8", "९": "9",
31 }
32 return "".join(hindi_digits.get(ch, ch) for ch in text)
33
34def extract_text_from_pdf(pdf_path: str, lang: str = "eng") -> str:
35 text = ""
36 poppler_path = get_poppler_path()
37
38 with open(pdf_path, "rb") as f:
39 reader = PyPDF2.PdfReader(f)
40 num_pages = len(reader.pages)
41
42 for page_num in range(num_pages):
43 page = reader.pages[page_num]
44 page_text = page.extract_text() or ""
45
46 if not page_text.strip():
47 try:
48 images = convert_from_path(
49 pdf_path,
50 first_page=page_num + 1,
51 last_page=page_num + 1,
52 poppler_path=(poppler_path if os.path.exists(os.path.join(poppler_path, "pdftoppm")) else None),
53 dpi=400,
54 grayscale=True,
55 )
56 if images:
57 img = images[0]
58 page_text = pytesseract.image_to_string(img, lang=lang, config="--psm 6")
59 except Exception as e:
60 st.error(f"OCR failed on page {page_num+1}: {e}")
61 page_text = ""
62
63 text += page_text + "\n"
64
65 return text
66
67def is_valid_page_number(page_str: str) -> bool:
68 if not page_str:
69 return False
70 return all(ch in "0123456789०१२३४५६७८९" for ch in page_str)
71
72def parse_toc(text: str, is_hindi: bool = False):
73 entries = []
74
75 common_patterns = [
76 r"^(.*?)[\s\.\-]+(\d+)\s*$",
77 r"^(.*?)[\s\-\_]+(\d+)\s*$",
78 r"^\s*(\d+\..*?)[\s\.\-]+(\d+)\s*$",
79 ]
80
81 hindi_patterns = [
82 r"^\s*([०१२३४५६७८९]+\.\s+.*?)[\s\.\-]*([०१२३४५६७८९\d]+)\s*$",
83 r"^(.*?)[\s\-\—]+([०१२३४५६७८९\d]+)\s*$",
84 r"^(.*?)[\s\.]+([०१२३४५६७८९\d]+)\s*$",
85 r"^(.*?(?:अध्याय|खंड|परिशिष्ट|प्रस्तावना|भाग|अनुभाग|प्रकरण)\s*[०१२३४५६७८९]*[\.\:\-]?\s*.*?)[\s\.\-]*([०१२३४५६७८९\d]+)\s*$",
86 r"^(.*?)\s+([०१२३४५६७८९\d]+)$",
87 ]
88
89 patterns = hindi_patterns if is_hindi else common_patterns
90 skip_terms_eng = ["table of contents", "contents", "page", "chap"]
91 skip_terms_hindi = ["विषय सूची", "अनुक्रमणिका", "सामग्री", "पृष्ठ", "अध्याय"]
92 skip_terms = skip_terms_hindi if is_hindi else skip_terms_eng
93
94 for line in text.split("\n"):
95 line = line.strip()
96 if not line or len(line) < 5:
97 continue
98 if any(term in line.lower() for term in skip_terms):
99 continue
100
101 for pattern in patterns:
102 match = re.match(pattern, line, re.IGNORECASE | re.UNICODE)
103 if not match:
104 continue
105
106 groups = match.groups()
107 if len(groups) == 2:
108 chapter = groups[0].strip()
109 page = groups[1].strip()
110 elif len(groups) == 3:
111 chapter = f"{groups[0]} {groups[1]}".strip()
112 page = groups[2].strip()
113 else:
114 continue
115
116 if is_valid_page_number(page):
117 page = convert_hindi_digits(page)
118 entries.append({"chapter": chapter, "page": page})
119 break
120 else:
121 fallback = re.search(r"(\d+|[०१२३४५६७८९]+)$", line)
122 if fallback:
123 page = fallback.group(1)
124 if is_valid_page_number(page):
125 chapter = line[: fallback.start()].strip()
126 page = convert_hindi_digits(page)
127 entries.append({"chapter": chapter, "page": page})
128 break
129
130 return entries
131
132
133# ─── Streamlit UI ──────────────────────────────────────────────────────────────
134
135st.set_page_config(page_title="PDF TOC Extractor", layout="centered")
136st.title("📄 PDF TOC Extractor")
137
138st.write(
139 """
140 1. Upload a PDF.
141 2. Use the Zoom slider to preview it.
142 3. Click “Extract TOC” → the app will parse and display it in an editable table.
143 4. Optionally add rows/columns at any index.
144 5. Finally, download a CSV named after your PDF.
145 """
146)
147
148# — Step 1: File Uploader & Zoom Slider —────────────────────────────────────────
149
150uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])
151zoom_pct = st.slider("Preview Zoom (%)", min_value=50, max_value=200, value=100, step=10)
152
153if uploaded_file:
154 # Display PDF preview inside an <iframe> with zoom scaled by zoom_pct
155 pdf_bytes = uploaded_file.read()
156 b64_pdf = base64.b64encode(pdf_bytes).decode("utf-8")
157 iframe_width = int(700 * (zoom_pct / 100))
158 iframe_height = 800
159
160 st.markdown(
161 f"""
162 <iframe
163 src="data:application/pdf;base64,{b64_pdf}"
164 width="{iframe_width}px"
165 height="{iframe_height}px"
166 style="border: none;"
167 ></iframe>
168 """,
169 unsafe_allow_html=True,
170 )
171
172 if "raw_pdf_bytes" not in st.session_state:
173 st.session_state["raw_pdf_bytes"] = pdf_bytes
174
175 # — Step 2: Extract TOC Button —───────────────────────────────────────────
176
177 language = st.selectbox("OCR Language (if needed)", ("eng", "hin", "both"))
178 if st.button("Extract TOC"):
179 with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
180 tmp.write(st.session_state["raw_pdf_bytes"])
181 tmp_path = tmp.name
182
183 is_hindi = (language == "hin" or language == "both")
184 ocr_lang = "eng+hin" if language == "both" else language
185
186 try:
187 raw_text = extract_text_from_pdf(tmp_path, lang=ocr_lang)
188 toc_list = parse_toc(raw_text, is_hindi=is_hindi)
189
190 if not toc_list:
191 st.warning("No TOC entries detected.")
192 else:
193 df = pd.DataFrame(toc_list)
194 st.session_state["df"] = df
195
196 except Exception as e:
197 st.error(f"Extraction error: {e}")
198 if os.path.exists(tmp_path):
199 os.remove(tmp_path)
200 st.stop()
201 finally:
202 if os.path.exists(tmp_path):
203 os.remove(tmp_path)
204
205# — Step 3: Editable Table & Add Row/Column —─────────────────────────────────
206
207if "df" in st.session_state:
208 st.subheader("🔧 Editable Table of Contents")
209 df = st.session_state["df"]
210
211 # Show the editable DataFrame (using st.data_editor instead of st.experimental_data_editor)
212 edited_df = st.data_editor(df, num_rows="dynamic", use_container_width=True)
213 st.session_state["df"] = edited_df
214
215 st.markdown("---")
216
217 # — Add a blank row at chosen index ───────────────────────────────────────
218 st.write("### ➕ Add a Blank Row")
219 max_row_idx = len(st.session_state["df"])
220 # Give the number_input its own key so we can read it on button click
221 st.number_input(
222 "Insert new row at index (0-based)",
223 min_value=0, max_value=max_row_idx, value=max_row_idx, step=1, key="new_row_idx"
224 )
225 if st.button("Add Row", key="add_row_button"):
226 df_current = st.session_state["df"]
227 new_row_idx = st.session_state["new_row_idx"]
228
229 # Create a one-row DataFrame of empty strings matching columns
230 blank_row = pd.DataFrame({col: [""] for col in df_current.columns})
231
232 # Split and concatenate at new_row_idx
233 top = df_current.iloc[: new_row_idx].reset_index(drop=True)
234 bottom = df_current.iloc[new_row_idx :].reset_index(drop=True)
235 new_df = pd.concat([top, blank_row, bottom], ignore_index=True)
236
237 st.session_state["df"] = new_df
238 st.experimental_rerun()
239
240 st.markdown("----")
241
242 # — Add a blank column at chosen index ──────────────────────────────────
243 st.write("### ➕ Add a Blank Column")
244 new_col_name = st.text_input("New column name", value="", key="new_col_name")
245 max_col_idx = len(st.session_state["df"].columns)
246 st.number_input(
247 "Insert new column at index (0-based)",
248 min_value=0, max_value=max_col_idx, value=max_col_idx, step=1, key="new_col_idx"
249 )
250 if st.button("Add Column", key="add_col_button"):
251 if st.session_state["new_col_name"].strip() == "":
252 st.error("Column name cannot be empty.")
253 else:
254 df_current = st.session_state["df"]
255 new_col_idx = st.session_state["new_col_idx"]
256 col_name = st.session_state["new_col_name"]
257
258 df_current.insert(new_col_idx, col_name, "")
259 st.session_state["df"] = df_current
260 st.experimental_rerun()
261
262 st.markdown("---")
263
264 # — Step 4: Download as CSV with PDF name ────────────────────────────────
265 st.write("### 💾 Download CSV")
266
267 original_pdf_name = uploaded_file.name
268 base_name = os.path.splitext(original_pdf_name)[0]
269 csv_filename = f"{base_name}.csv"
270
271 final_df = st.session_state["df"]
272 csv_data = final_df.to_csv(index=False).encode("utf-8")
273
274 st.download_button(
275 label="Download TOC as CSV",
276 data=csv_data,
277 file_name=csv_filename,
278 mime="text/csv",
279 )
280