app.py

import cv2
import numpy as np
import json
import gradio as gr
import os
import xml.etree.ElementTree as ET

# ---------------- Helper functions ----------------
def get_rotated_rect_corners(x, y, w, h, rotation_deg):
    rot_rad = np.deg2rad(rotation_deg)
    cos_r = np.cos(rot_rad)
    sin_r = np.sin(rot_rad)
    R = np.array([[cos_r, -sin_r],
                  [sin_r,  cos_r]])
    cx = x + w/2
    cy = y + h/2
    local_corners = np.array([
        [-w/2, -h/2],
        [ w/2, -h/2],
        [ w/2,  h/2],
        [-w/2,  h/2]
    ])
    rotated_corners = np.dot(local_corners, R.T)
    corners = rotated_corners + np.array([cx, cy])
    return corners.astype(np.float32)

def preprocess_gray_clahe(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    return clahe.apply(gray)

def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
    if method == "SIFT":
        detector = cv2.SIFT_create(nfeatures=5000)
        norm = cv2.NORM_L2
    elif method == "ORB":
        detector = cv2.ORB_create(5000)
        norm = cv2.NORM_HAMMING
    elif method == "BRISK":
        detector = cv2.BRISK_create()
        norm = cv2.NORM_HAMMING
    elif method == "KAZE":
        detector = cv2.KAZE_create()
        norm = cv2.NORM_L2
    elif method == "AKAZE":
        detector = cv2.AKAZE_create()
        norm = cv2.NORM_HAMMING
    else:
        return None, None, [], None

    kp1, des1 = detector.detectAndCompute(img1_gray, None)
    kp2, des2 = detector.detectAndCompute(img2_gray, None)
    if des1 is None or des2 is None:
        return None, None, [], None

    matcher = cv2.BFMatcher(norm)
    raw_matches = matcher.knnMatch(des1, des2, k=2)
    good = [m for m,n in raw_matches if m.distance < ratio_thresh * n.distance]

    matches_img = None
    if len(good) >= 4:
        matches_img = cv2.drawMatches(
            cv2.cvtColor(img1_gray, cv2.COLOR_GRAY2BGR),
            kp1,
            cv2.cvtColor(img2_gray, cv2.COLOR_GRAY2BGR),
            kp2,
            good, None,
            flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS
        )
    return kp1, kp2, good, matches_img

def add_title(img_bgr, title):
    h, w = img_bgr.shape[:2]
    bar = np.full((40, w, 3), 255, dtype=np.uint8)
    cv2.putText(bar, title, (10, 28), cv2.FONT_HERSHEY_SIMPLEX,
                0.8, (0,0,0), 2, cv2.LINE_AA)
    return np.vstack([bar, img_bgr])

def parse_xml_roi_points(xml_path):
    """Parse your XML structure, return list of polygons (Nx2 points)."""
    polys = []
    try:
        tree = ET.parse(xml_path)
        root = tree.getroot()
        # Transform ROI points (FourPoint)
        for tr in root.findall(".//transform"):
            pts = []
            for p in tr.findall("point"):
                x = float(p.get("x")); y = float(p.get("y"))
                pts.append([x, y])
            if len(pts) >= 3:
                polys.append(np.array(pts, dtype=np.float32))
        # Overlay polygons
        for ov in root.findall(".//overlay"):
            pts = []
            for p in ov.findall("point"):
                x = float(p.get("x")); y = float(p.get("y"))
                pts.append([x, y])
            if len(pts) >= 3:
                polys.append(np.array(pts, dtype=np.float32))
    except Exception as e:
        print("XML parse error:", e)
    return polys if polys else None

# ---------------- Main Function ----------------
def homography_all_detectors(flat_file, persp_file, json_file, xml_file):
    flat_img = cv2.imread(flat_file)
    persp_img = cv2.imread(persp_file)
    mockup = json.load(open(json_file.name))
    roi = mockup["printAreas"][0]
    roi_x, roi_y = roi["position"]["x"], roi["position"]["y"]
    roi_w, roi_h = roi["width"], roi["height"]
    roi_rot_deg = roi["rotation"]

    xml_polys = parse_xml_roi_points(xml_file.name) if xml_file else None

    flat_gray = preprocess_gray_clahe(flat_img)
    persp_gray = preprocess_gray_clahe(persp_img)

    results = []
    download_files = []

    for method in ["SIFT","ORB","BRISK","KAZE","AKAZE"]:
        kp1, kp2, good, matches_img = detect_and_match(flat_gray, persp_gray, method)
        if kp1 is None or len(good) < 4:
            continue

        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1,1,2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1,1,2)
        H, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

        # top-left = flat
        top_left = add_title(flat_img.copy(), "Flat Image")

        # top-right = feature matches
        top_right = add_title(matches_img if matches_img is not None else flat_img.copy(),
                              "Feature Matches (Flat→Perspective)")

        # bottom-left = homography ROI
        bottom_left = persp_img.copy()
        if H is not None:
            roi_corners = get_rotated_rect_corners(roi_x, roi_y, roi_w, roi_h, roi_rot_deg)
            roi_persp = cv2.perspectiveTransform(roi_corners.reshape(-1,1,2), H).reshape(-1,2)
            cv2.polylines(bottom_left, [roi_persp.astype(int)], True, (0,255,0), 3)
        bottom_left = add_title(bottom_left, "ROI via Homography (from JSON)")

        # bottom-right = XML ROI
        bottom_right = persp_img.copy()
        if xml_polys:
            for poly in xml_polys:
                cv2.polylines(bottom_right, [poly.astype(int)], True, (0,0,255), 3)
        else:
            cv2.putText(bottom_right, "No XML ROI", (50,100),
                        cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255), 3)
        bottom_right = add_title(bottom_right, "Ground Truth ROI (XML)")

        # stack horizontally
        top = np.hstack([top_left, top_right])
        bottom = np.hstack([bottom_left, bottom_right])
        composite = np.vstack([top, bottom])

        base_name = os.path.splitext(os.path.basename(persp_file))[0]
        file_name = f"{base_name}_{method.lower()}_grid.png"
        cv2.imwrite(file_name, composite)
        results.append((cv2.cvtColor(composite, cv2.COLOR_BGR2RGB), f"{method} Grid"))
        download_files.append(file_name)

    while len(download_files)<5: download_files.append(None)
    return [results]+download_files[:5]

# ---------------- Gradio UI ----------------
iface = gr.Interface(
    fn=homography_all_detectors,
    inputs=[
        gr.Image(label="Upload Flat Image", type="filepath"),
        gr.Image(label="Upload Perspective Image", type="filepath"),
        gr.File(label="Upload mockup.json", file_types=[".json"]),
        gr.File(label="Upload Groundtruth.xml", file_types=[".xml"])
    ],
    outputs=[
        gr.Gallery(label="Composite Grids (per Detector)", show_label=True),
        gr.File(label="Download SIFT Grid"),
        gr.File(label="Download ORB Grid"),
        gr.File(label="Download BRISK Grid"),
        gr.File(label="Download KAZE Grid"),
        gr.File(label="Download AKAZE Grid")
    ],
    title="Homography ROI vs XML ROI",
    description="Each detector produces one 2×2 grid: Flat, Matches, Homography ROI, Ground Truth ROI."
)
iface.launch()