import gradio as gr
from model import load_models, run_inference

FACE_WEIGHTS = "weights/face_detector_fasterrcnn_final.pth"
LANDMARK_WEIGHTS = "weights/landmark_keypointrcnn_final.pth"

NUM_KEYPOINTS = 12  # change if needed

face_model, landmark_model = load_models(
    FACE_WEIGHTS,
    LANDMARK_WEIGHTS,
    NUM_KEYPOINTS
)


def predict(image):
    return run_inference(image, face_model, landmark_model)


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    title="Cattle Face Landmark Detection",
    description="Upload an image of cattle. The model detects the face and predicts facial landmarks."
)

demo.launch()