from fastai.vision.all import *
import gradio as gr
import cv2
classifier = cv2.CascadeClassifier('haarcascade_frontalface_alt2.xml')

def label_func(fname):
    if int(str(fname)[str(fname).index('_')+1]) == 0:
        return "Male"
    return "Female"

def get_age(fname):
    return int(str(fname).split('/')[1].split('_')[0])

def detect_face(img):
    faces = classifier.detectMultiScale(img)
    x, y, w, h = faces[0]
    cropped_img = img[y:y+h, x:x+w]
    return cropped_img

learn_gender = load_learner('gender.pkl')
learn_age = load_learner('age.pkl')

categories = ('Female', 'Male')

def predict_age(img):
    detected_face = detect_face(img)
    pred,_,_ = learn_age.predict(detected_face)
    return str(pred[0]), detected_face

def classify_image(img):
    pred, idx, probs = learn_gender.predict(img)
    return dict(zip(categories, map(float, probs)))

def process_image(img):
    gender = classify_image(img)
    age, face = predict_age(img)
    return gender, age, face

image = gr.inputs.Image(shape=(192,192))
gender_output = gr.outputs.Label()
age_output = gr.outputs.Textbox(label='Predicted Age')
detected_face_output = gr.outputs.Image(type='numpy', label='Detected Face')
examples = ['Male.jpg', 'Female.png']


iface = gr.Interface(fn=process_image, inputs=image, outputs=[gender_output, age_output, detected_face_output], examples=examples)
iface.launch()