import cv2
import numpy as np
import tensorflow as tf
from tensorflow.keras.utils import custom_object_scope

class CustomMSE(tf.keras.losses.Loss):
    def __init__(self, name="custom_mse"):
        super().__init__(name=name)

    def call(self, y_true, y_pred):
        return tf.reduce_mean(tf.square(y_true - y_pred), axis=-1)

def predict(image):
    
    
    model = tf.keras.models.load_model("v1model.h5")
    
    img = cv2.imread(image)
    
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    haar_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

    faces_rect = haar_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=9)
    men_count = 0
    women_count = 0
    for (x, y, w, h) in faces_rect:
        cropped_img = img[y:y+h, x:x+w]

        resized_img = cv2.resize(cropped_img,(64,64))
        image_array = np.expand_dims(resized_img,axis=0)

        prediction = model.predict(image_array)


        if prediction[0][0] > 0.5:
            ans = "Men"
            men_count+=1
        else:
            ans = "Women"
            women_count+=1
    

        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), thickness=2)
        cv2.putText(img, ans, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return rgb_img,men_count,women_count