import gradio as gr
import cv2
import numpy as np
import json

def detect_aruco_markers(image, dictionary_preset,
                         cornerRefinementMethod, detectInvertedMarker, adaptiveThreshConstant):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Load predefined dictionary
    aruco_dict = cv2.aruco.getPredefinedDictionary(getattr(cv2.aruco, dictionary_preset))

    # Initialize detector parameters
    parameters = cv2.aruco.DetectorParameters()
    parameters.cornerRefinementMethod = getattr(cv2.aruco, cornerRefinementMethod)
    parameters.detectInvertedMarker = detectInvertedMarker
    parameters.adaptiveThreshConstant = adaptiveThreshConstant

    detector = cv2.aruco.ArucoDetector(aruco_dict, parameters)

    # Detect markers
    corners, ids, rejected = detector.detectMarkers(gray)
    print("Detected IDs:", ids)

    detections = []
    if ids is not None:
        annotated_image = cv2.aruco.drawDetectedMarkers(image.copy(), corners, ids)
        for i in range(len(ids)):
            marker_info = {
                "id": int(ids[i][0]),
                "corners": [np.array(corner).tolist() for corner in corners[i]]
            }
            detections.append(marker_info)
    else:
        annotated_image = image.copy()

    return annotated_image, json.dumps(detections, indent=4)

# Available ArUco dictionary presets
dictionary_choices = [name for name in dir(cv2.aruco) if name.startswith("DICT_")]
corner_refinement_choices = ["CORNER_REFINE_NONE", "CORNER_REFINE_SUBPIX", "CORNER_REFINE_CONTOUR", "CORNER_REFINE_APRILTAG"]

# Gradio interface
iface = gr.Interface(
    fn=detect_aruco_markers,
    inputs=[
        gr.Image(type="numpy"),
        gr.Dropdown(choices=dictionary_choices, label="Dictionary Preset", value="DICT_6X6_250"),	
        gr.Dropdown(choices=corner_refinement_choices, label="Corner Refinement Method", value="CORNER_REFINE_NONE"),
        gr.Checkbox(label="Detect Inverted Marker", value=False),
        gr.Slider(minimum=3, maximum=201, step=2, label="Adaptive Thresh", value=7),
    ],
    outputs=[
        gr.Image(type="numpy", label="Detected Markers"),
        gr.Code(label="Detection Data (JSON)", language="json")
    ],
    title="ArUco Marker Detector with Parameter Control",
    description="Upload an image and adjust ArUco detection parameters. The detected markers are drawn on the image, and the detection data (IDs and corner coordinates) is provided in JSON format.",
)

if __name__ == "__main__":
    iface.launch()