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