import cv2
import gradio as gr
import numpy as np
from PIL import Image

article = "References<br>[1] Z. Shen et al., LayoutParser: A Unified Toolkit for Deep Learning Based Document Image Analysis, 2021."
description = "Upload a document image to detect prominent layout blocks. This CPU-friendly version uses image processing instead of downloading a brittle external EfficientDet checkpoint at startup."


def show_preds(input_image):
    if input_image is None:
        return None

    image = np.array(input_image.convert("RGB"))
    resized = image.copy()
    max_width = 1200
    if resized.shape[1] > max_width:
        scale = max_width / resized.shape[1]
        resized = cv2.resize(resized, (max_width, int(resized.shape[0] * scale)))

    gray = cv2.cvtColor(resized, cv2.COLOR_RGB2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 35, 11)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 10))
    closed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    annotated = resized.copy()
    min_area = max(1000, int(0.002 * resized.shape[0] * resized.shape[1]))
    boxes = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w * h >= min_area and w > 30 and h > 20:
            boxes.append((x, y, w, h))

    for idx, (x, y, w, h) in enumerate(sorted(boxes, key=lambda b: (b[1], b[0])), start=1):
        cv2.rectangle(annotated, (x, y), (x + w, y + h), (255, 90, 0), 3)
        cv2.putText(annotated, f"Block {idx}", (x, max(20, y - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 90, 0), 2)

    return Image.fromarray(annotated)


demo = gr.Interface(
    fn=show_preds,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Document Layout Detector/Parser",
    article=article,
    description=description,
    analytics_enabled=True,
)
demo.queue().launch()