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, sin_r = np.cos(rot_rad), np.sin(rot_rad)
    R = np.array([[cos_r, -sin_r], [sin_r, cos_r]])
    cx, cy = x + w/2, 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)
    return (rotated_corners + np.array([cx,cy])).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); matcher=cv2.BFMatcher(cv2.NORM_L2)
    elif method=="ORB": detector=cv2.ORB_create(5000); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
    elif method=="BRISK": detector=cv2.BRISK_create(); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
    elif method=="KAZE": detector=cv2.KAZE_create(); matcher=cv2.BFMatcher(cv2.NORM_L2)
    elif method=="AKAZE": detector=cv2.AKAZE_create(); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
    else: return 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,[]
    
    raw_matches = matcher.knnMatch(des1,des2,k=2)
    good = [m for m,n in raw_matches if m.distance < ratio_thresh*n.distance]
    return kp1, kp2, good

def parse_xml_points(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    transform = root.find('.//transform')
    points = {}
    for pt in transform.findall('.//point'):
        pt_type = pt.attrib['type']
        x = float(pt.attrib['x'])
        y = float(pt.attrib['y'])
        points[pt_type] = (x, y)
    return points

# ---------------- Fit-to-Box Helper ----------------
def fit_to_box(img, target_h=600, target_w=600):
    h, w = img.shape[:2]
    scale = min(target_w/w, target_h/h)   # preserve aspect ratio
    new_w, new_h = int(w*scale), int(h*scale)
    resized = cv2.resize(img, (new_w, new_h))

    # symmetric padding
    top = (target_h - new_h) // 2
    bottom = target_h - new_h - top
    left = (target_w - new_w) // 2
    right = target_w - new_w - left

    canvas = np.ones((target_h, target_w, 3), dtype=np.uint8) * 255
    canvas[top:top+new_h, left:left+new_w] = resized
    return canvas

# ---------------- NEW: Remap keypoints to boxed image coords ----------------
def remap_keypoints_to_box(kps, orig_shape, target_h=600, target_w=600):
    h, w = orig_shape[:2]
    scale = min(target_w / w, target_h / h)
    new_w, new_h = int(w * scale), int(h * scale)
    top = (target_h - new_h) // 2
    left = (target_w - new_w) // 2

    kps_new = []
    for kp in kps:
        x, y = kp.pt
        x_new = x * scale + left
        y_new = y * scale + top
        kps_new.append(cv2.KeyPoint(x_new, y_new, max(1.0, kp.size * scale),
                                    kp.angle, kp.response, kp.octave, kp.class_id))
    return kps_new

# ---------------- Add Heading on Top ----------------
def add_heading(img, text):
    # add white band on top
    h, w = img.shape[:2]
    band_h = 40
    canvas = np.ones((h+band_h, w, 3), dtype=np.uint8) * 255
    canvas[band_h:] = img
    cv2.putText(canvas, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                1, (0,0,0), 2, cv2.LINE_AA)
    return canvas

# ---------------- 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_data = mockup["printAreas"][0]["position"]
    roi_x, roi_y = roi_data["x"], roi_data["y"]
    roi_w, roi_h = mockup["printAreas"][0]["width"], mockup["printAreas"][0]["height"]
    roi_rot_deg = mockup["printAreas"][0]["rotation"]

    flat_gray = preprocess_gray_clahe(flat_img)
    persp_gray = preprocess_gray_clahe(persp_img)
    xml_points = parse_xml_points(xml_file.name)

    methods = ["SIFT","ORB","BRISK","KAZE","AKAZE"]
    gallery_paths = []
    download_files = []

    for method in methods:
        kp1,kp2,good_matches = detect_and_match(flat_gray,persp_gray,method)
        if kp1 is None or kp2 is None or len(good_matches)<4: continue

        flat_box = fit_to_box(flat_img, 600, 600)
        persp_box = fit_to_box(persp_img, 600, 600)

        kp1_box = remap_keypoints_to_box(kp1, flat_img.shape, 600, 600)
        kp2_box = remap_keypoints_to_box(kp2, persp_img.shape, 600, 600)

        match_img = cv2.drawMatches(flat_box, kp1_box, persp_box, kp2_box, good_matches, None, flags=2)

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

        roi_corners_flat = get_rotated_rect_corners(roi_x,roi_y,roi_w,roi_h,roi_rot_deg)
        roi_corners_persp = cv2.perspectiveTransform(roi_corners_flat.reshape(-1,1,2),H).reshape(-1,2)
        persp_roi = persp_img.copy()
        cv2.polylines(persp_roi,[roi_corners_persp.astype(int)],True,(0,255,0),2)
        for px,py in roi_corners_persp: cv2.circle(persp_roi,(int(px),int(py)),5,(255,0,0),-1)

        xml_gt_img = persp_img.copy()
        ordered_pts = ['TopLeft', 'TopRight', 'BottomRight', 'BottomLeft']
        xml_polygon = [xml_points[pt] for pt in ordered_pts]
        pts = np.array(xml_polygon, np.int32).reshape((-1,1,2))
        cv2.polylines(xml_gt_img,[pts],isClosed=True,color=(255,0,0),thickness=3)

        flat_rgb = fit_to_box(cv2.cvtColor(flat_img,cv2.COLOR_BGR2RGB),600,600)
        match_rgb = fit_to_box(cv2.cvtColor(match_img,cv2.COLOR_BGR2RGB),600,600)
        roi_rgb = fit_to_box(cv2.cvtColor(persp_roi,cv2.COLOR_BGR2RGB),600,600)
        xml_rgb = fit_to_box(cv2.cvtColor(xml_gt_img,cv2.COLOR_BGR2RGB),600,600)

        # Add headings
        flat_rgb = add_heading(flat_rgb, "Flat Image")
        match_rgb = add_heading(match_rgb, "Flat -> Perspective Feature Matching")
        roi_rgb = add_heading(roi_rgb, "Perspective Image with Homography ROI  ")
        xml_rgb = add_heading(xml_rgb, "    Perspective GT ROI")

        # Merge 2x2 grid
        top = np.hstack([flat_rgb, match_rgb])
        bottom = np.hstack([roi_rgb, xml_rgb])
        combined_grid = np.vstack([top, bottom])

        base_name = os.path.splitext(os.path.basename(persp_file))[0]
        file_name = f"{base_name}_{method.lower()}.png"
        cv2.imwrite(file_name, cv2.cvtColor(combined_grid,cv2.COLOR_RGB2BGR))
        gallery_paths.append(file_name)
        download_files.append(file_name)

    while len(download_files)<5: download_files.append(None)
    return gallery_paths, download_files[0], download_files[1], download_files[2], download_files[3], download_files[4]

# ---------------- 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 XML file",file_types=[".xml"])
    ],
    outputs=[
        gr.Gallery(label="Results per Detector",show_label=True),
        gr.File(label="Download SIFT Result"),
        gr.File(label="Download ORB Result"),
        gr.File(label="Download BRISK Result"),
        gr.File(label="Download KAZE Result"),
        gr.File(label="Download AKAZE Result")
    ],
    title="Homography ROI + Feature Matching + XML GT",
    description="Flat + Perspective images with mockup.json & XML. Aspect ratio preserved, images centered in uniform boxes, headings added."
)

iface.launch()