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

MODEL_PATH = "best.pt"
model = YOLO(MODEL_PATH)

def draw_pose_clean(image, conf=0.25, show_labels=False, point_size=3):
    if image is None:
        return None

    img = np.array(image.convert("RGB"))
    results = model.predict(source=img, conf=conf, imgsz=640, verbose=False)[0]

    out = img.copy()
    if results.keypoints is None or len(results.keypoints.xy) == 0:
        return Image.fromarray(out)

    kpts_xy = results.keypoints.xy.cpu().numpy()
    kpts_conf = results.keypoints.conf.cpu().numpy() if results.keypoints.conf is not None else None

    for det_i, det_kpts in enumerate(kpts_xy):
        confs = kpts_conf[det_i] if kpts_conf is not None else np.ones(len(det_kpts))
        for i, (x, y) in enumerate(det_kpts):
            score = float(confs[i])
            if score < conf:
                continue

            x_i, y_i = int(x), int(y)
            cv2.circle(out, (x_i, y_i), int(point_size), (0, 255, 0), -1, lineType=cv2.LINE_AA)

            # Optional: only label every 5th keypoint to reduce clutter
            if show_labels and (i % 5 == 0):
                cv2.putText(
                    out,
                    f"{i+1}",
                    (x_i + 4, y_i - 4),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.35,
                    (255, 255, 0),
                    1,
                    cv2.LINE_AA,
                )

    return Image.fromarray(out)

demo = gr.Interface(
    fn=draw_pose_clean,
    inputs=[
        gr.Image(type="pil", label="Input image"),
        gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence threshold"),
        gr.Checkbox(value=False, label="Show keypoint index labels (sparser)"),
        gr.Slider(1, 8, value=3, step=1, label="Point size"),
    ],
    outputs=gr.Image(type="pil", label="Output"),
    title="DogFLW Pose (Clean View)",
    description="Less cluttered output: points only by default. Toggle sparse labels if needed.",
)

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

# from PIL import Image
# import numpy as np
# import gradio as gr
# from ultralytics import YOLO

# MODEL_PATH = "best.pt"
# model = YOLO(MODEL_PATH)

# def predict(image, conf):
#     if image is None:
#         return None
#     results = model.predict(source=np.array(image), conf=conf, imgsz=640, verbose=False)[0]
#     plotted = results.plot()  # BGR numpy array
#     plotted = plotted[:, :, ::-1]  # BGR -> RGB
#     return Image.fromarray(plotted)

# demo = gr.Interface(
#     fn=predict,
#     inputs=[
#         gr.Image(type="pil", label="Input Image"),
#         gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence"),
#     ],
#     outputs=gr.Image(type="pil", label="Pose Result"),
#     title="DogFLW YOLOv8 Pose",
#     description="Upload a dog image to detect 46 facial landmarks."
# )

# if __name__ == "__main__":
#     demo.launch()