import tensorflow as tf
import cv2
import numpy as np
# import matplotlib.pyplot as plt
# import seaborn as sns
import streamlit as st
from PIL import Image

st.image(r"face-mask.png", width = 150)
st.write("# Mask detector App")

# load model
@st.cache_resource
def cache_model(model_add):
    model = tf.keras.models.load_model(model_add)
    return model

model = cache_model("mask_detector.keras")

labels = ['WithMask', 'WithoutMask']
def predict(img,model,labels = labels):
  # read_image
#   img = cv2.imread(img_add)
  # resize image to 224, 224
  img = cv2.resize(img, (224,224))
  # st.write("resize_cv2", img.shape)
  # convert to RGB
  img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  img_expand = tf.expand_dims(img_rgb, axis = 0)
  y_pred_prob = model.predict(img_expand)
  y_pred_prob_round = np.round(y_pred_prob[0][0], 3)
  st.write(y_pred_prob[0][0])
  y_pred = int(np.round(y_pred_prob_round))
  label = labels[y_pred]
  return (label, y_pred_prob_round)

def face_detect(img, model, labels = labels):
#   img = cv2.imread(img_add)
  face_detector=cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
  img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  faces = face_detector.detectMultiScale(img_gray)
  if len(faces) > 0:
    for face in faces:
      x,y,w,h = face
      x_max = x + w
      y_max = y + h
      face_crp = img[y:y_max+1, x:x_max+1]
      label, prob = predict(face_crp,model,labels = labels)
      if label == labels[0]:
        cv2.rectangle(img, (x,y), (x_max, y_max), (0,255,0), 3)
        cv2.putText(img, f"{label} -: {prob}%", (x,y-10), cv2.FONT_HERSHEY_SIMPLEX,
                     1, (0,255,0), 2)
      else:
        cv2.rectangle(img, (x,y), (x_max, y_max), (255,0,0), 3)
        cv2.putText(img, f"{label} -: {1 - prob}%", (x,y-10), cv2.FONT_HERSHEY_SIMPLEX,
                     1, (255,0,0), 2)
    return img


# File uploader
uploaded_file = st.file_uploader("upload image", type=["jpg", "jpeg", "png"])
submit_btn = st.button("Submit")

if submit_btn:
  if uploaded_file:
    img_array = np.array(Image.open(uploaded_file))
    # st.write(img_array.shape)
    result_img = face_detect(img_array, model)
    # result_img_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
    st.image(result_img)
  else:
    st.write("Please upload an image")