import gradio as gr
import cv2
import numpy as np
from paddleocr import PPStructureV3 # Explicitly import the class that exists

# --- INITIALIZATION ---
# We do NOT pass a custom model path. We let PPStructureV3 download its own default model.
# This avoids the "ValueError: Unknown argument" crashes.
layout_engine = PPStructureV3(
    use_doc_orientation_classify=True, # Standard V3 argument for orientation
    enable_mkldnn=False                # CRITICAL: Keeps CPU from crashing
)

def analyze_layout(input_image):
    if input_image is None:
        return None, "No image uploaded"

    image_np = np.array(input_image)
    
    # Run Inference
    try:
        # V3 returns a generator, so we convert to list immediately
        results = list(layout_engine(image_np))
    except Exception as e:
        return image_np, f"Error running layout analysis: {e}"

    viz_image = image_np.copy()
    detections_text = []

    if not results:
        return viz_image, "No layout detected."
        
    # --- VISUALIZATION ---
    for region in results:
        if not isinstance(region, dict): continue
        
        # V3 usually puts the box in 'layout_bbox' or 'bbox'
        box = region.get('layout_bbox') or region.get('bbox')
        label = region.get('label', 'unknown')
        
        if box is None: continue
        
        try:
            x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
            
            # Color coding
            color = (0, 255, 0) # Green (Default)
            if label == 'title': color = (0, 0, 255)    # Red
            elif label == 'figure': color = (255, 0, 0) # Blue
            elif label == 'table': color = (255, 255, 0)# Cyan

            cv2.rectangle(viz_image, (x1, y1), (x2, y2), color, 3)
            cv2.putText(viz_image, str(label), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
            detections_text.append(f"Found {label} at {box}")
        except Exception:
            pass

    return viz_image, "\n".join(detections_text)

with gr.Blocks(title="PP-DocLayoutV3 Explorer") as demo:
    gr.Markdown("## 📄 PP-DocLayoutV3 Explorer")
    gr.Markdown("Auto-downloading the latest V3 weights for structure analysis.")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input Document")
            submit_btn = gr.Button("Analyze Layout", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Layout Visualization")
            output_log = gr.Textbox(label="Detected Regions", lines=10)

    submit_btn.click(fn=analyze_layout, inputs=input_img, outputs=[output_img, output_log])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)