import cv2
import gradio as gr
from gradio_point_promptable_image import PointPromptableImage

# Define the colors to use for marking points
BLUE = (135, 206, 235)
PINK = (239, 149, 186)


# Function to extract specific point inputs based on criteria
def get_point_inputs(prompts):
    point_inputs = []
    for prompt in prompts:
        # Example condition to filter points (you can adjust this)
        if prompt[5] == 4.0:  # Checks for a specific condition in the prompts
            point_inputs.append((prompt[0], prompt[1], prompt[2]))
    return point_inputs


# Function to process the image and mark points based on user input
def process_input(input_dict):
    img, points = input_dict["image"], input_dict["points"]

    # Extract points from user input
    point_inputs = get_point_inputs(points)

    # Mark the points on the image
    for point in point_inputs:
        x, y = int(point[0]), int(point[1])
        cv2.circle(
            img, (x, y), 2, (0, 0, 0), thickness=10
        )  # Black border for visibility
        if point[2] == 1:
            cv2.circle(img, (x, y), 2, BLUE, thickness=8)  # Blue for class 1
        else:
            cv2.circle(img, (x, y), 2, PINK, thickness=8)  # Pink for other classes

    return img  # Return the image with the points marked


# Create the Gradio interface with an uploadable image
demo = gr.Interface(
    process_input,  # The function to process the input image and points
    PointPromptableImage(),  # Custom input component to get points from the user
    gr.Image(),  # Output component to display the processed image
    title="Image Upload and Point Marker",
    description="Upload an image, click on it to mark points, and see the points marked.",
)

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()