app.py

import gradio as gr
from PIL import Image, ImageDraw
import requests
from io import BytesIO
import numpy as np
import json
import tempfile
import easyocr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from bs4 import BeautifulSoup
import base64
import re

# ----------------- Initialize OCR -----------------
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
reader = easyocr.Reader(['en'])

# ----------------- HTML Parsing -----------------
from bs4 import BeautifulSoup

from bs4 import BeautifulSoup

def parse_html_to_json(html_file):
    """
    Parse HTML content from a Gradio file input or string and produce
    words/paragraphs JSON compatible with image OCR output.
    """
    # Handle Gradio NamedString, str, or file-like object
    html_content = ""
    if hasattr(html_file, "read"):  # real file
        html_content = html_file.read()
        if isinstance(html_content, bytes):
            html_content = html_content.decode("utf-8")
    elif isinstance(html_file, str):
        html_content = html_file
    else:  # Gradio NamedString
        html_content = getattr(html_file, "name", str(html_file))

    soup = BeautifulSoup(html_content, "html.parser")

    words_json = []
    paragraphs_json = []
    y_offset = 0
    line_height = 20
    char_width = 10

    # iterate over all visible text nodes in the body
    body = soup.body
    if not body:
        body = soup  # fallback

    # Only consider visible text
    for element in body.find_all(text=True):
        text = element.strip()
        if not text:
            continue

        # split into words
        line_words = text.split()
        line_bbox = [0, y_offset, char_width * len(text), y_offset + line_height]

        word_entries = []
        x_offset = 0
        for word in line_words:
            word_bbox = [x_offset, y_offset, x_offset + char_width * len(word), y_offset + line_height]
            word_entry = {"text": word, "bbox": word_bbox, "confidence": 1.0}
            word_entries.append(word_entry)
            words_json.append(word_entry)
            x_offset += char_width * (len(word) + 1)

        paragraphs_json.append({
            "text": text,
            "bbox": line_bbox,
            "words": word_entries
        })

        y_offset += line_height

    output_json = {
        "words": words_json,
        "paragraphs": paragraphs_json
    }

    return output_json

# ----------------- Image Loading -----------------
def load_image(image_file, image_url):
    if image_file:
        return [image_file]
    elif image_url:
        response = requests.get(image_url)
        return [Image.open(BytesIO(response.content)).convert("RGB")]
    return []

# ----------------- Main Logic -----------------
def detect_text_combined(image_file, image_url, html_file):
    # ----------------- HTML Path -----------------
    if html_file:
        output_json = parse_html_to_json(html_file)
        json_str = json.dumps(output_json, indent=2)
        tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w")
        tmp_file.write(json_str)
        tmp_file.close()
        annotated_image = None
        return annotated_image, json_str, tmp_file.name

    # ----------------- Image Path -----------------
    images = load_image(image_file, image_url)
    if not images:
        return None, "No input provided.", None

    annotated_image = images[0]
    image = annotated_image
    results = reader.readtext(np.array(image))
    draw = ImageDraw.Draw(image)
    words_json = []

    for bbox, _, conf in results:
        x_coords = [float(point[0]) for point in bbox]
        y_coords = [float(point[1]) for point in bbox]
        x_min, y_min = min(x_coords), min(y_coords)
        x_max, y_max = max(x_coords), max(y_coords)

        # Crop word for TrOCR recognition
        word_crop = image.crop((x_min, y_min, x_max, y_max))
        pixel_values = processor(images=word_crop, return_tensors="pt").pixel_values
        generated_ids = model.generate(pixel_values)
        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=2)

        words_json.append({
            "text": text,
            "bbox": [x_min, y_min, x_max, y_max],
            "confidence": float(conf)
        })

    paragraphs_json = words_json.copy()
    output_json = {
        "words": words_json,
        "paragraphs": paragraphs_json
    }
    json_str = json.dumps(output_json, indent=2)
    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w")
    tmp_file.write(json_str)
    tmp_file.close()

    return annotated_image, json_str, tmp_file.name

# ----------------- Gradio Interface -----------------
iface = gr.Interface(
    fn=detect_text_combined,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(label="Image URL (optional)"),
        gr.File(label="Upload HTML File", file_types=[".html", ".htm"])
    ],
    outputs=[
        gr.Image(type="pil", label="Annotated Image"),
        gr.Textbox(label="JSON Output"),
        gr.File(label="Download JSON")
    ],
    title="Combined OCR & HTML Text Bounding Box Extractor",
    description="Upload an image, provide an image URL, or upload an HTML file. Outputs word- and paragraph-level bounding boxes in JSON format consistent with image OCR output."
)

if __name__ == "__main__":
    iface.launch()