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()