app.py

import os
import io
import cv2
import gradio as gr
import pandas as pd
import easyocr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from lxml import etree
from ultralytics import YOLO
import tempfile

# Load Arial with bigger size (e.g., 24)
font = ImageFont.truetype("arial.ttf", size=24)

# === Load models ===
model = YOLO("model/NER_66_37_50_V2.pt")  # Update path if needed
reader = easyocr.Reader(['fr'])

# Global store
stored_detections = []
stored_image = None
image_name = "uploaded_image.jpg"

# === Helper Functions ===
def detect_yolo(image_path):
    global stored_detections, stored_image, image_name

    image_name = os.path.basename(image_path)  # <-- Extract just filename
    image_pil = Image.open(image_path).convert("RGB")
    # image_pil = image_pil.convert("RGB")
    stored_image = np.array(image_pil)
    results = model(stored_image)[0]
    draw = ImageDraw.Draw(image_pil)
    detections = []

    for box in results.boxes.data.tolist():
        coords = list(map(int, box[:4]))
        conf = float(box[4])
        cls = int(box[5])
        x1, y1, x2, y2 = coords
        class_name = model.names[cls]

        draw.rectangle([(x1, y1), (x2, y2)], outline="red", width=6)
        draw.text((x1, y1 - 30), class_name, fill="blue", font=font)

        detections.append({
            "image": image_name,
            "class": class_name,
            "x1": x1, "y1": y1,
            "x2": x2, "y2": y2,
            "yolo_conf": round(conf, 4),
            "text": "", "ocr_conf": 0.0
        })

    stored_detections = detections
    return image_pil

def recognize_text():
    global stored_detections, stored_image

    for det in stored_detections:
        x1, y1, x2, y2 = det['x1'], det['y1'], det['x2'], det['y2']
        crop = stored_image[y1:y2, x1:x2]
        ocr_result = reader.readtext(crop, detail=1)

        if ocr_result:
            det['text'] = " ".join([res[1] for res in ocr_result])
            det['ocr_conf'] = round(sum(res[2] for res in ocr_result) / len(ocr_result), 4)
        else:
            det['text'], det['ocr_conf'] = "", 0.0

    summary = "\n".join([f"[{d['class']}] {d['text']} (conf: {d['ocr_conf']})" for d in stored_detections])
    return summary

# def generate_csv():
#     df = pd.DataFrame(stored_detections)
#     csv_buffer = io.StringIO()
#     df.to_csv(csv_buffer, index=False)
#     csv_buffer.seek(0)
#     # return df, gr.File.update(value=csv_buffer, filename="results.csv")
#     return df, "results.csv"

def generate_csv():
    df = pd.DataFrame(stored_detections)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", mode='w', encoding='utf-8', newline='')
    df.to_csv(temp_file.name, index=False)
    temp_file.close()
    # return df, temp_file.name
    csv_filename = os.path.splitext(image_name)[0] + ".csv"
    temp_file_path = os.path.join(tempfile.gettempdir(), csv_filename)
    df.to_csv(temp_file_path, index=False)
    return df, temp_file_path

def generate_alto_xml():
    H, W, _ = stored_image.shape
    NSMAP = {"alto": "http://www.loc.gov/standards/alto/ns-v4#"}
    root = etree.Element("alto", nsmap=NSMAP)

    desc = etree.SubElement(root, "Description")
    etree.SubElement(desc, "MeasurementUnit").text = "pixel"
    etree.SubElement(etree.SubElement(desc, "sourceImageInformation"), "fileName").text = image_name

    layout = etree.SubElement(root, "Layout")
    page = etree.SubElement(layout, "Page", ID="page_1", PHYSICAL_IMG_NR="1", HEIGHT=str(H), WIDTH=str(W))
    ps = etree.SubElement(page, "PrintSpace", HEIGHT=str(H), WIDTH=str(W), HPOS="0", VPOS="0")

    for i, d in enumerate(stored_detections):
        w, h = d['x2'] - d['x1'], d['y2'] - d['y1']
        block = etree.SubElement(ps, "TextBlock", ID=f"b_{i}", HPOS=str(d['x1']), VPOS=str(d['y1']),
                                 WIDTH=str(w), HEIGHT=str(h))
        line = etree.SubElement(block, "TextLine", ID=f"l_{i}", HPOS=str(d['x1']), VPOS=str(d['y1']),
                                WIDTH=str(w), HEIGHT=str(h))
        etree.SubElement(line, "String", CONTENT=d['text'], HPOS=str(d['x1']), VPOS=str(d['y1']),
                         WIDTH=str(w), HEIGHT=str(h),
                         WC=str(d['ocr_conf']), CLASS=d['class'], YOLO_CONF=str(d['yolo_conf']))

    xml_bytes = etree.tostring(root, pretty_print=True, xml_declaration=True, encoding="UTF-8")

    xml_text = xml_bytes.decode("utf-8")

    # buffer = io.BytesIO(xml_bytes)
    # buffer.seek(0)
    # # return xml_text, gr.File.update(value=buffer, filename="result.xml")

    # with open("result.xml", "wb") as f:
    #     f.write(xml_bytes)

    # return xml_text, "result.xml"

    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".xml", mode="wb")
    temp_file.write(xml_bytes)
    temp_file.close()

    # return xml_text, temp_file.name
    xml_filename = os.path.splitext(image_name)[0] + ".xml"
    temp_file_path = os.path.join(tempfile.gettempdir(), xml_filename)

    with open(temp_file_path, "wb") as f:
        f.write(xml_bytes)

    return xml_text, temp_file_path

# === Gradio UI Layout ===
with gr.Blocks(title="TextVision AI: From Detection to XML") as demo:
    gr.Markdown("# 🧾 TextVision AI: From Detection to XML")

    # input_img = gr.Image(label="📁 Upload Image", type="pil")
    input_img = gr.Image(label="📁 Upload Image", type="filepath")

    # === Load Sample Images ===

    example_images = [
        ["example/Image_1.jpg"],
        ["example/Image_2.jpg"],
        ["example/Image_3.jpg"]
    ]

    gr.Examples(
    examples=example_images,
    inputs=input_img,
    label="🖼️ Example Images"
    )

    btn_detect = gr.Button("🔲 Step 1: Detect")
    detected_img = gr.Image(label="🖼️ Detected Image")

    btn_ocr = gr.Button("🔠 Step 2: Recognize")
    ocr_output = gr.Textbox(label="📝 Recognized Text with Tags", lines=6)

    btn_csv = gr.Button("📄 Step 3: Generate CSV")
    csv_table = gr.Dataframe(label="📊 Detection Table")
    csv_file = gr.File(label="⬇️ Download CSV")

    btn_xml = gr.Button("🗂️ Step 4: Generate ALTO XML")
    xml_text = gr.Textbox(label="📄 ALTO XML View", lines=12)
    xml_file = gr.File(label="⬇️ Download XML")

    # Bind buttons
    btn_detect.click(fn=detect_yolo, inputs=input_img, outputs=detected_img)
    btn_ocr.click(fn=recognize_text, outputs=ocr_output)
    btn_csv.click(fn=generate_csv, outputs=[csv_table, csv_file])
    btn_xml.click(fn=generate_alto_xml, outputs=[xml_text, xml_file])

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