import gradio as gr
import cv2, numpy as np
from PIL import Image
from fastai.learner import load_learner

learn = load_learner('best_model.pkl')

# use the _alt2 cascade and lower the minNeighbors
face_cascade = cv2.CascadeClassifier(
    cv2.data.haarcascades + "haarcascade_frontalface_alt2.xml"
)

def predict(img: Image.Image):
    # convert to array & detect
    arr = np.array(img.convert("RGB"))
    gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)
    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.05,    # smaller steps
        minNeighbors=3,      # accept weaker detections
        minSize=(30,30)
    )

    if len(faces)==0:
        # fallback: center-crop a square region
        w,h = img.size
        side = min(w,h)
        left = (w-side)//2; top = (h-side)//2
        face_img = img.crop((left, top, left+side, top+side))
    else:
        x,y,w,h = faces[0]
        face_img = img.crop((x,y,x+w,y+h))

    # run your model
    pred,_,probs = learn.predict(face_img)
    p = probs.max().item()
    if p<0.6: return {"The image could not be categorised. Please try with another photo.":1.0}
    return {str(c): float(probs[i]) for i,c in enumerate(learn.dls.vocab)}



iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", interactive=True),
    outputs=gr.Label(num_top_classes=3),
    title="Age Category Classifier",
    description="Upload a face image and get probabilities for young / middle / old."
)

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