Update model and draw boundary

app.py

@@ -10,110 +10,25 @@ import PIL
 import io
 import html
 import time
-
+from facenet_pytorch import MTCNN
 
 model_file_path = 'models/bmi.h5'
 model = tf.keras.models.load_model(model_file_path)
 
+mtcnn2 = MTCNN(
+        image_size=160, margin=40, min_face_size=20,
+        thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=False
+    )
+
 # Preprocess the images for VGG16
 def preprocess_image(img_path):
     img = image.load_img(img_path, target_size = (224, 224))
     img = image.img_to_array(img)
     img = np.expand_dims(img, axis = 0)
-    img = preprocess_input(img)
+    img = preprocess_input(img)/255.
     return img
 
 
-# function to convert OpenCV Rectangle bounding box image into base64 byte string to be overlayed on video stream
-def bbox_to_bytes(bbox_array):
-  """
-  Params:
-          bbox_array: Numpy array (pixels) containing rectangle to overlay on video stream.
-  Returns:
-        bytes: Base64 image byte string
-  """
-  # convert array into PIL image
-  bbox_PIL = PIL.Image.fromarray(bbox_array, 'RGBA')
-  iobuf = io.BytesIO()
-  # format bbox into png for return
-  bbox_PIL.save(iobuf, format='png')
-  # format return string
-  bbox_bytes = 'data:image/png;base64,{}'.format((str(b64encode(iobuf.getvalue()), 'utf-8')))
-
-  return bbox_bytes
-
-# base_model = VGGFace(model='vgg16', include_top=False, input_shape=(224, 224, 3))
-# x = base_model.output
-# x = GlobalAveragePooling2D()(x)
-# model = Model(inputs=base_model.input, outputs=x)
-
-# # Function to preprocess the image
-# def preprocess_image(img):
-#     img = cv2.resize(img, (224, 224))
-#     img = image.img_to_array(img)
-#     img = np.expand_dims(img, axis=0)
-#     img = img[0] # Remove the extra dimension
-#     return img
-    
-# def extract_features(image_array):
-#     # img = np.squeeze(image_array, axis=0)
-#     img = np.expand_dims(image_array, axis=0)
-#     img = tf.keras.applications.resnet50.preprocess_input(img)
-#     features = model.predict(img,verbose=0)
-#     return features.flatten()
-
-# Function to predict BMI
-
-def draw_boundary(img):
-    # initialize the Haar Cascade face detection model
-    face_cascade = cv2.CascadeClassifier(cv2.samples.findFile(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'))
-
-    # initialze bounding box to empty
-    bbox = ''
-    count = 0 
-    while True:
-
-        # create transparent overlay for bounding box
-        bbox_array = np.zeros([480,640,4], dtype=np.uint8)
-
-        # grayscale image for face detection
-        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-
-        # get face region coordinates
-        faces = face_cascade.detectMultiScale(gray)
-        # get face bounding box for overlay
-        for (x, y, w, h) in faces:
-            # Extract the face region from the frame
-            face = img[y:y+h, x:x+w]
-            
-            # Preprocess the face image
-            face = cv2.resize(face, (224, 224))
-            face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
-            face = preprocess_input(face)/255.
-            face = np.expand_dims(face, axis=0)
-            
-            # Predict BMI using the pre-trained model
-            bmi = model.predict(face)[0][0]
-            
-            # Draw the predicted BMI on the frame
-            bbox_array = cv2.putText(bbox_array, f'BMI: {bmi:.2f}', (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-            
-            # Draw a rectangle around the face
-            bbox_array = cv2.rectangle(bbox_array, (x, y), (x+w, y+h), (255, 0, 0), 2)
-
-        bbox_array[:,:,3] = (bbox_array.max(axis = 2) > 0 ).astype(int) * 255
-        # convert overlay of bbox into bytes
-        bbox_bytes = bbox_to_bytes(bbox_array)
-        # update bbox so next frame gets new overlay
-        bbox = bbox_bytes
-    
-    return img
-
-def predict_bmi(img):
-    pre_img = preprocess_image(img)
-    pred = draw_boundary(pre_img)
-    return pred
-
 def main():
     st.title("BMI Prediction from Camera Image")
     st.write("This app predicts the BMI of a person from an image captured using the camera.")
@@ -122,15 +37,15 @@ def main():
     img_file_buffer = st.camera_input("Take a picture")
 
     if img_file_buffer is not None:
-        # Load the image data from the file buffer
-        file_bytes = np.asarray(bytearray(img_file_buffer.getvalue()), dtype=np.uint8)
-        img = cv2.imdecode(file_bytes, 1)
-
-        # Preprocess the image and predict BMI
-        bmi_label = predict_bmi(img)
-
-        # Display the predicted BMI
-        st.write("Predicted BMI:", str(bmi_label[0] - 5))
+        # To read image file buffer as a PIL Image:
+        img = Image.open(img_file_buffer)
+
+        detected_image = Image.fromarray(mtcnn2(img).numpy().transpose(1, 2, 0).astype(np.uint8))
+        st.image(detected_image, caption="Detected Face")
+        
+        embeddings = preprocess_image(img_file_buffer)
+        bmi = round(model.predict(embeddings), 2) - 4
+        st.write(f"Your BMI is {bmi}")
 
 if __name__ == '__main__':
     main()