import streamlit as st
from PIL import Image, ImageDraw
import os
import io
from backend.text_editor.pdf_processor import get_pdf_page_count, pdf_page_to_pil, compile_images_to_pdf
from backend.text_editor.ocr_engine import EASYOCR_AVAILABLE
from backend.text_editor.orchestrator import TextEditorOrchestrator
from backend.text_editor.font_detector import load_matching_font
def load_demo_invoice():
"""Create a beautiful synthetic demo invoice image to test instantly."""
w, h = 700, 900
img = Image.new("RGB", (w, h), "#FAFAFA")
draw = ImageDraw.Draw(img)
draw.rectangle([20, 20, w-20, h-20], fill="#FFFFFF", outline="#E0E0E0", width=2)
try:
from PIL import ImageFont
font_large = ImageFont.truetype("C:\\Windows\\Fonts\\arial.ttf", 36)
font_med = ImageFont.truetype("C:\\Windows\\Fonts\\arial.ttf", 16)
font_bold = ImageFont.truetype("C:\\Windows\\Fonts\\arialbd.ttf", 16)
except Exception:
font_large = ImageFont.load_default()
font_med = ImageFont.load_default()
font_bold = ImageFont.load_default()
draw.text((50, 60), "INVOICE", fill="#1A1B35", font=font_large)
draw.text((50, 140), "Invoice Number: INV-2026-9042", fill="#555555", font=font_med)
draw.text((50, 170), "Date: May 29, 2026", fill="#555555", font=font_med)
draw.text((50, 200), "Billed To: Antigravity User", fill="#555555", font=font_med)
draw.line([50, 240, w-50, 240], fill="#E0E0E0", width=1)
draw.text((50, 260), "Description", fill="#1A1B35", font=font_bold)
draw.text((450, 260), "Quantity", fill="#1A1B35", font=font_bold)
draw.text((550, 260), "Amount", fill="#1A1B35", font=font_bold)
draw.line([50, 290, w-50, 290], fill="#1A1B35", width=2)
draw.text((50, 310), "AI Background Remover Software Licence", fill="#555555", font=font_med)
draw.text((450, 310), "1", fill="#555555", font=font_med)
draw.text((550, 310), "$149.00", fill="#555555", font=font_med)
draw.text((50, 350), "Advanced Guided Filter Matting Module", fill="#555555", font=font_med)
draw.text((450, 350), "2", fill="#555555", font=font_med)
draw.text((550, 350), "$199.00", fill="#555555", font=font_med)
draw.line([50, 420, w-50, 420], fill="#E0E0E0", width=1)
draw.text((400, 450), "Subtotal:", fill="#555555", font=font_med)
draw.text((550, 450), "$348.00", fill="#555555", font=font_med)
draw.text((400, 480), "Tax (10%):", fill="#555555", font=font_med)
draw.text((550, 480), "$34.80", fill="#555555", font=font_med)
draw.text((400, 520), "Total Due:", fill="#1A1B35", font=font_bold)
draw.text((550, 520), "$382.80", fill="#1A1B35", font=font_bold)
draw.rectangle([50, 780, w-50, 840], fill="#F0F4F8", outline="#D0DCE5")
draw.text((70, 795), "Thank you for your business! For queries contact support@antigravity.ai", fill="#4A5D6E", font=font_med)
st.session_state.te_original_file = img
st.session_state.te_file_type = "image"
st.session_state.te_pdf_filename = "demo_invoice.png"
st.session_state.te_original_image = img
st.session_state.te_pdf_pages = [img]
st.session_state.te_edited_image = None
st.session_state.te_replacements = []
st.session_state.te_ocr_results = None
def draw_ocr_overlays_with_badges(pil_img: Image.Image, ocr_results: list) -> Image.Image:
"""Overlays high-quality translucent bounding boxes and numbered circular badges over text areas."""
overlay_img = pil_img.copy()
draw = ImageDraw.Draw(overlay_img, "RGBA")
try:
from PIL import ImageFont
font_badge = ImageFont.truetype("C:\\Windows\\Fonts\\arial.ttf", 13)
except Exception:
font_badge = ImageFont.load_default()
for idx, item in enumerate(ocr_results):
x, y, w, h = item["bbox"]
# Draw translucent cyan bounding box
draw.rectangle([x, y, x + w, y + h], fill=(0, 240, 255, 15), outline=(0, 240, 255, 160), width=1)
# Draw small solid pink badge at top-left corner of bounding box
bx, by = x - 5, y - 5
br = 9 # radius
draw.ellipse([bx - br, by - br, bx + br, by + br], fill=(255, 0, 127, 240))
# Center number in circle
num_str = str(idx + 1)
draw.text((bx - 4, by - 6), num_str, fill=(255, 255, 255, 255), font=font_badge)
return overlay_img
def generate_font_preview_pil(text: str, font_name: str) -> Image.Image:
"""Generates a small visual PIL image showing the replacement text rendered in the Google Font."""
w, h = 280, 50
canvas = Image.new("RGB", (w, h), "#F8F9FA")
draw = ImageDraw.Draw(canvas)
try:
font, font_size = load_matching_font(bbox_height=26, font_family=font_name, size_multiplier=1.0, text=text)
draw.text((10, 12), text, fill=(26, 27, 53), font=font)
except Exception:
draw.text((10, 15), text, fill=(26, 27, 53))
return canvas
def render_text_editor_ui():
"""Renders the modular AI In-Image & PDF Document Text Editor workspace with interactive numbered badges."""
st.markdown(
"""
📝 AI Document & PDF Text Editor
Correct text directly using visual numbered badges. Search and preview the most matching Google Fonts.
""",
unsafe_allow_html=True
)
# ----------------- SESSION STATE INITIALIZATION -----------------
if "te_original_file" not in st.session_state:
st.session_state.te_original_file = None
if "te_file_type" not in st.session_state:
st.session_state.te_file_type = None
if "te_pdf_filename" not in st.session_state:
st.session_state.te_pdf_filename = "document"
if "te_original_image" not in st.session_state:
st.session_state.te_original_image = None
if "te_pdf_pages" not in st.session_state:
st.session_state.te_pdf_pages = []
if "te_active_page" not in st.session_state:
st.session_state.te_active_page = 0
if "te_ocr_results" not in st.session_state:
st.session_state.te_ocr_results = None
if "te_edited_image" not in st.session_state:
st.session_state.te_edited_image = None
if "te_replacements" not in st.session_state:
st.session_state.te_replacements = []
# ----------------- SIDEBAR CONTROLS -----------------
st.sidebar.markdown("### ⚙️ TEXT EDITOR CONTROLS")
font_family = st.sidebar.selectbox(
"Default Font Family",
["Sans-Serif", "Serif", "Monospace"],
index=0,
help="Selects the nearest matching system TTF font style for text rendering."
)
custom_font = st.sidebar.text_input(
"⚡ Or Search Google Font",
value="",
placeholder="e.g. Rozha One, Kalam, Hind, Poppins",
help="Type any Google Font family name. The engine will dynamically download, extract, and cache the full font set!"
)
if custom_font.strip():
font_family = custom_font.strip()
size_multiplier = st.sidebar.slider(
"Font Size Scale Adjust",
min_value=0.5,
max_value=1.5,
value=0.90,
step=0.05,
help="Scale multiplier to perfectly size the drawn characters to the surrounding bounding box."
)
overlay_ocr = st.sidebar.checkbox(
"Overlay Numbered Badges",
value=True,
help="Draws glowing pink circular badges and cyan boxes over text blocks to make editing easy."
)
# ----------------- MAIN WORKFLOW -----------------
# Case 1: No file loaded yet
if st.session_state.te_original_file is None:
col_l, col_c, col_r = st.columns([1, 3, 1])
with col_c:
st.markdown(
"""
⚡ UPLOAD SOURCE DOCUMENT
Upload an Image or multi-page PDF Document. Coordinates and fonts will be parsed instantly.
""",
unsafe_allow_html=True
)
uploaded_file = st.file_uploader(
"Choose document or image...",
type=["pdf", "jpg", "jpeg", "png", "webp"],
label_visibility="collapsed",
key="doc_uploader"
)
if uploaded_file is not None:
try:
file_bytes = uploaded_file.read()
filename = uploaded_file.name
st.session_state.te_pdf_filename = filename
if filename.lower().endswith(".pdf"):
st.session_state.te_original_file = file_bytes
st.session_state.te_file_type = "pdf"
page_count = get_pdf_page_count(file_bytes)
with st.spinner(f"Extracting {page_count} PDF pages at high quality (150 DPI)..."):
pdf_pages = []
for idx in range(page_count):
pdf_pages.append(pdf_page_to_pil(file_bytes, idx, dpi=150))
st.session_state.te_pdf_pages = pdf_pages
st.session_state.te_active_page = 0
st.session_state.te_original_image = pdf_pages[0]
else:
from PIL import ImageOps
pil_img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
pil_img = ImageOps.exif_transpose(pil_img)
st.session_state.te_original_file = file_bytes
st.session_state.te_file_type = "image"
st.session_state.te_pdf_pages = [pil_img]
st.session_state.te_active_page = 0
st.session_state.te_original_image = pil_img
st.session_state.te_edited_image = None
st.session_state.te_replacements = []
st.session_state.te_ocr_results = None
st.rerun()
except Exception as e:
st.error(f"Error loading document: {str(e)}")
st.markdown("No files ready? Test instantly with a pre-configured demo invoice:
", unsafe_allow_html=True)
col_demo_left, col_demo_center, col_demo_right = st.columns([1, 1, 1])
with col_demo_center:
if st.button("Load Demo Invoice", use_container_width=True):
load_demo_invoice()
st.rerun()
# Case 2: Document loaded
else:
col_nav, col_reset = st.columns([5, 1])
with col_reset:
if st.button("Reset / Clear", key="te_reset", use_container_width=True):
st.session_state.te_original_file = None
st.session_state.te_file_type = None
st.session_state.te_pdf_filename = "document"
st.session_state.te_original_image = None
st.session_state.te_pdf_pages = []
st.session_state.te_ocr_results = None
st.session_state.te_edited_image = None
st.session_state.te_replacements = []
st.rerun()
with col_nav:
if st.session_state.te_file_type == "pdf":
pages_count = len(st.session_state.te_pdf_pages)
st.markdown(
f"""
📄 File: {st.session_state.te_pdf_filename} | Total Pages: {pages_count}
""",
unsafe_allow_html=True
)
if pages_count > 1:
active_pg = st.slider("Navigate Pages", 1, pages_count, st.session_state.te_active_page + 1) - 1
if active_pg != st.session_state.te_active_page:
st.session_state.te_active_page = active_pg
st.session_state.te_original_image = st.session_state.te_pdf_pages[active_pg]
st.session_state.te_ocr_results = None
st.session_state.te_edited_image = None
st.session_state.te_replacements = []
st.rerun()
else:
st.markdown(
f"""
🖼️ File: {st.session_state.te_pdf_filename} ({st.session_state.te_original_image.width}x{st.session_state.te_original_image.height}px)
""",
unsafe_allow_html=True
)
# ----------------- OCR BACKGROUND TASK -----------------
if EASYOCR_AVAILABLE and st.session_state.te_ocr_results is None:
with st.spinner("Analyzing image text layout using OCR engine..."):
ocr_res = TextEditorOrchestrator.run_ocr(st.session_state.te_original_image)
st.session_state.te_ocr_results = ocr_res
# ----------------- DUAL COLUMN LAYOUT -----------------
col_preview, col_editor_panel = st.columns([1, 1])
# Left Column: Image Previews (Before/After)
with col_preview:
st.markdown('1. ORIGINAL SOURCE DOCUMENT
', unsafe_allow_html=True)
# Display image with circular number badges
if overlay_ocr and st.session_state.te_ocr_results:
original_badged = draw_ocr_overlays_with_badges(st.session_state.te_original_image, st.session_state.te_ocr_results)
st.image(original_badged, use_container_width=True)
else:
st.image(st.session_state.te_original_image, use_container_width=True)
st.markdown('2. EDITED DOCUMENT PREVIEW
', unsafe_allow_html=True)
if st.session_state.te_edited_image is not None:
st.image(st.session_state.te_edited_image, use_container_width=True)
else:
st.markdown(
"""
No edits made yet.
Modify a text block on the right panel to see updates.
""",
unsafe_allow_html=True
)
# Right Column: Visual In-Place Editor Panel
with col_editor_panel:
st.markdown('3. CLICK & EDIT DOCUMENT TEXT FIELDS
', unsafe_allow_html=True)
if not st.session_state.te_ocr_results:
st.warning("⚠️ No text detected in this page. Try uploading an image with higher contrast or clear characters.")
else:
# Present list of all detected blocks as visual numbered rows
st.markdown('', unsafe_allow_html=True)
# Dictionary to store user's changed inputs
user_replacements = {}
for idx, item in enumerate(st.session_state.te_ocr_results):
# Bounding Box description row
orig_text = item["text"]
bbox = item["bbox"]
st.markdown(
f"""
{idx+1}
Original:
"{orig_text}"
""",
unsafe_allow_html=True
)
# Edit input box
rep_input = st.text_input(
"Replacement Value",
value=orig_text,
key=f"rep_val_{idx}",
label_visibility="collapsed"
)
# Store if the user actually edited it
if rep_input != orig_text:
user_replacements[idx] = {
"bbox": bbox,
"original_text": orig_text,
"replacement_text": rep_input
}
st.markdown('
4. MATCHING GOOGLE FONT PREVIEW GALLERY
', unsafe_allow_html=True)
st.info("💡 **Google Font Search**: Choose a text block on the left and see how it renders in these popular Google Fonts in real-time, then type the matching name into **Or Search Google Font** in the sidebar!")
# Let user pick which replacement string to preview in the gallery
words_to_preview = [item["text"] for item in st.session_state.te_ocr_results]
selected_preview_word = st.selectbox("Select word to preview fonts:", words_to_preview)
# Font Showcase grid
st.markdown('', unsafe_allow_html=True)
hindi_showcase_fonts = {
"Rozha One": "Thick Display Serifs (Perfect for headers)",
"Noto Sans Devanagari": "Sleek Modern Sans-Serif",
"Yatra One": "Classic Indian Calligraphy",
"Kalam": "Handwritten/Artistic Calligraphy",
"Hind": "Clean corporate Sans-Serif",
"Martel": "Formal Document Serif",
"Teko": "Condensed Bold Display"
}
# Display 3 popular fonts previews side-by-side
col_f1, col_f2 = st.columns(2)
font_list = list(hindi_showcase_fonts.keys())
for i, font_name in enumerate(font_list):
col_target = col_f1 if i % 2 == 0 else col_f2
with col_target:
st.markdown(f"**{font_name}** (*{hindi_showcase_fonts[font_name]}*)")
preview_pil = generate_font_preview_pil(selected_preview_word, font_name)
st.image(preview_pil, use_container_width=True)
st.markdown("
", unsafe_allow_html=True)
st.markdown('
💾 EXPORT DOCUMENT
""",
unsafe_allow_html=True
)
st.markdown('', unsafe_allow_html=True)
filename_base = st.session_state.te_pdf_filename.rsplit('.', 1)[0]
if st.session_state.te_file_type == "pdf":
col_dl_pdf, col_dl_cur = st.columns(2)
with col_dl_pdf:
with st.spinner("Compiling all multi-page PDF pages..."):
compiled_pdf_bytes = compile_images_to_pdf(st.session_state.te_pdf_pages)
st.markdown(
"""
📄
Complete Edited PDF
Re-compiles all pages into a single PDF document.
""",
unsafe_allow_html=True
)
st.download_button(
label="Download Full Document PDF",
data=compiled_pdf_bytes,
file_name=f"{filename_base}_edited.pdf",
mime="application/pdf",
key="dl_edited_pdf",
use_container_width=True
)
with col_dl_cur:
pg_buf = io.BytesIO()
st.session_state.te_edited_image.save(pg_buf, format="PNG")
pg_bytes = pg_buf.getvalue()
st.markdown(
"""
🖼️
Active Page Image (PNG)
Downloads the current page as a single high-DPI image.
""",
unsafe_allow_html=True
)
st.download_button(
label=f"Download Page {st.session_state.te_active_page+1} PNG",
data=pg_bytes,
file_name=f"{filename_base}_page_{st.session_state.te_active_page+1}_edited.png",
mime="image/png",
key="dl_edited_page_png",
use_container_width=True
)
else:
img_buf = io.BytesIO()
st.session_state.te_edited_image.save(img_buf, format="PNG")
img_bytes = img_buf.getvalue()
col_dl_img_l, col_dl_img_c, col_dl_img_r = st.columns([1, 2, 1])
with col_dl_img_c:
st.markdown(
"""
✨
Edited Image Output
Downloads the text-corrected image as high-fidelity PNG.
""",
unsafe_allow_html=True
)
st.download_button(
label="Download Edited Image PNG",
data=img_bytes,
file_name=f"{filename_base}_edited.png",
mime="image/png",
key="dl_edited_img_png",
use_container_width=True
)
st.markdown('