app.py

import gradio as gr
import cv2
import numpy as np
import os
import json
from PIL import Image
import io
import base64
from openai import OpenAI
from ultralytics import YOLO

# Define the Latex2Layout model path
model_path = "latex2layout_object_detection_yolov8.pt"

# Verify model file existence
if not os.path.exists(model_path):
    raise FileNotFoundError(f"Model file not found at {model_path}")

# Load the Latex2Layout model with error handling
try:
    model = YOLO(model_path)
except Exception as e:
    raise RuntimeError(f"Failed to load Latex2Layout model: {e}")

# Qwen API configuration
QWEN_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
QWEN_MODELS = {
    "Qwen2.5-VL-3B-Instruct": "qwen2.5-vl-3b-instruct",
    "Qwen2.5-VL-7B-Instruct": "qwen2.5-vl-7b-instruct",
    "Qwen2.5-VL-14B-Instruct": "qwen2.5-vl-14b-instruct",
}

# Default system prompt template
default_system_prompt = """You are an assistant specialized in document layout analysis.
The following layout elements were detected in the image (confidence >= 0.5):
{layout_info}
Use this information and the image to answer layout-related questions."""

def encode_image(image_array):
    """
    Convert a numpy array image to a base64-encoded string.

    Args:
        image_array: Numpy array representing the image.

    Returns:
        str: Base64-encoded string of the image.
    """
    try:
        pil_image = Image.fromarray(image_array)
        img_byte_arr = io.BytesIO()
        pil_image.save(img_byte_arr, format='PNG')
        return base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
    except Exception as e:
        raise ValueError(f"Failed to encode image: {e}")

def detect_layout(image, confidence_threshold=0.5):
    """
    Detect layout elements in the uploaded image using the Latex2Layout model.

    Args:
        image: Uploaded image as a numpy array.
        confidence_threshold: Minimum confidence score to retain detections (default: 0.5).

    Returns:
        tuple: (annotated_image, layout_info_str)
            - annotated_image: Image with bounding boxes drawn (confidence >= 0.5).
            - layout_info_str: JSON string of layout detections (confidence >= 0.5).
    """
    if image is None or not isinstance(image, np.ndarray):
        return None, "Error: No image uploaded or invalid image format."

    try:
        # Perform detection
        results = model(image)
        result = results[0]
        annotated_image = image.copy()
        layout_info = []

        # Process detections
        for box in result.boxes:
            conf = float(box.conf[0])
            if conf < confidence_threshold:
                continue

            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            cls_id = int(box.cls[0])
            cls_name = result.names[cls_id]

            color = tuple(np.random.randint(0, 255, 3).tolist())
            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)
            label = f"{cls_name} {conf:.2f}"
            (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(annotated_image, (x1, y1 - label_height - 5), (x1 + label_width, y1), color, -1)
            cv2.putText(annotated_image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

            layout_info.append({
                "bbox": [x1, y1, x2, y2],
                "class": cls_name,
                "confidence": conf
            })

        layout_info_str = json.dumps(layout_info, indent=2) if layout_info else "No layout elements detected with confidence >= 0.5."
        return annotated_image, layout_info_str

    except Exception as e:
        return None, f"Error during layout detection: {str(e)}"

def detect_example_image():
    """
    Load and detect layout elements in the example image (./image1.png).

    Returns:
        tuple: (example_image, annotated_image, layout_info_str)
            - example_image: Original example image.
            - annotated_image: Annotated example image.
            - layout_info_str: JSON string of layout detections.
    """
    example_image_path = "./image1.png"
    if not os.path.exists(example_image_path):
        return None, None, "Error: Example image not found."

    try:
        # Load image in BGR and convert to RGB
        bgr_image = cv2.imread(example_image_path)
        if bgr_image is None:
            return None, None, "Error: Failed to load example image."
        rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)

        # Run detection
        annotated_image, layout_info_str = detect_layout(rgb_image)
        return rgb_image, annotated_image, layout_info_str
    except Exception as e:
        return None, None, f"Error processing example image: {str(e)}"

def qa_about_layout(image, question, layout_info, api_key, model_name, system_prompt_template):
    """
    Answer layout-related questions using the Qwen API with an editable system prompt.

    Args:
        image: Uploaded image as a numpy array.
        question: User's question about the layout.
        layout_info: JSON string of layout detection results.
        api_key: User's Qwen API key.
        model_name: Selected Qwen model name.
        system_prompt_template: Editable system prompt template.

    Returns:
        str: Qwen's response to the question.
    """
    if image is None or not isinstance(image, np.ndarray):
        return "Error: Please upload a valid image."
    if not question:
        return "Error: Please enter a question."
    if not api_key:
        return "Error: Please provide a Qwen API key."
    if not layout_info:
        return "Error: No layout information available. Detect layout first."

    try:
        # Encode image to base64
        base64_image = encode_image(image)

        # Map model name to ID
        model_id = QWEN_MODELS.get(model_name)
        if not model_id:
            return "Error: Invalid Qwen model selected."

        # Replace placeholder in system prompt with layout info
        system_prompt = system_prompt_template.replace("{layout_info}", layout_info)

        # Initialize OpenAI client for Qwen API
        client = OpenAI(api_key=api_key, base_url=QWEN_BASE_URL)

        # Prepare API request messages
        messages = [
            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}},
                    {"type": "text", "text": question},
                ],
            },
        ]

        # Call Qwen API
        completion = client.chat.completions.create(model=model_id, messages=messages)
        return completion.choices[0].message.content

    except Exception as e:
        return f"Error during QA: {str(e)}"

# Build Gradio interface
with gr.Blocks(title="Latex2Layout QA System") as demo:
    gr.Markdown("# Latex2Layout QA System")
    gr.Markdown("Upload an image or use the example to detect layout elements and ask questions using Qwen models.")

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Upload Image", type="numpy")
            detect_btn = gr.Button("Detect Layout")
            example_btn = gr.Button("Detect Example Image")
            gr.Markdown("**Tip**: Use clear images for best results.")

        with gr.Column(scale=1):
            output_image = gr.Image(label="Detected Layout")
            layout_info = gr.Textbox(label="Layout Information", lines=10, interactive=False)

    with gr.Row():
        with gr.Column(scale=1):
            api_key_input = gr.Textbox(
                label="Qwen API Key",
                placeholder="Enter your Qwen API key",
                type="password"
            )
            model_select = gr.Dropdown(
                label="Select Qwen Model",
                choices=list(QWEN_MODELS.keys()),
                value="Qwen2.5-VL-3B-Instruct"
            )
            gr.Markdown("**System Prompt Template**: Edit the prompt sent to Qwen. Include `{layout_info}` to insert detection results.")
            system_prompt_input = gr.Textbox(
                label="System Prompt Template",
                value=default_system_prompt,
                lines=5,
                placeholder="Edit the system prompt here. Keep {layout_info} to include detection results."
            )
            question_input = gr.Textbox(label="Ask About the Layout", placeholder="e.g., 'Where is the heading?'")
            qa_btn = gr.Button("Ask Question")

        with gr.Column(scale=1):
            answer_output = gr.Textbox(label="Answer", lines=5, interactive=False)

    # Event handlers
    detect_btn.click(
        fn=detect_layout,
        inputs=[input_image],
        outputs=[output_image, layout_info]
    )
    example_btn.click(
        fn=detect_example_image,
        inputs=[],
        outputs=[input_image, output_image, layout_info]
    )
    qa_btn.click(
        fn=qa_about_layout,
        inputs=[input_image, question_input, layout_info, api_key_input, model_select, system_prompt_input],
        outputs=[answer_output]
    )

# Launch the application
if __name__ == "__main__":
    demo.launch()