Create classifier.py

classifier.py

@@ -0,0 +1,88 @@
+import cv2
+import os
+import numpy as np
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.models import load_model
+from tensorflow.keras.models import model_from_json
+import time
+
+
+class Classifier:
+    def __init__(self):
+        self.is_image = None
+        self.classifier = None
+        self.face_detector = None
+        self.labels = {0: "Angry", 1: "Disgusted", 2: "Fearful",
+                       3: "Happy", 4: "Neutral", 5: "Surprised", 6: "Sad"}
+        self.load_models()
+    
+    def check_file_type(self, is_image, input_path):
+        self.is_image = is_image
+        self.input_path = input_path
+        self.output_path = "result/" + self.input_path.split('/')[-1]
+        os.makedirs("result", exist_ok=True)
+
+    def load_models(self):
+        # Load VGG
+        yaml_file = open('models/VGG19.yaml', 'r')
+        self.classifier = model_from_json(yaml_file.read())
+        self.classifier.load_weights("models/VGG19.h5")
+        yaml_file.close()
+        print("VGG19 has been loaded")
+
+        # Load Face Detector
+        self.face_detector = cv2.CascadeClassifier(
+        'models/haarcascade_frontalface_default.xml')
+        cv2.ocl.setUseOpenCL(False)
+        print("Face Detector has been loaded")
+
+    def predict(self, frame):
+
+        faces = self.face_detector.detectMultiScale(
+            frame, scaleFactor=1.3, minNeighbors=5)
+
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        for (x, y, w, h) in faces:
+            cv2.rectangle(frame, (x, y-50), (x+w, y+h+10), (255, 0, 0), 2)
+            roi_gray = gray[y:y + h, x:x + w]
+            roi_color = frame[y:y + h, x:x + w]
+            # Convert grayscale to RGB
+            roi_color_rgb = cv2.cvtColor(roi_color, cv2.COLOR_BGR2RGB)
+            cropped_img = np.expand_dims(np.expand_dims(
+                cv2.resize(roi_color_rgb, (48, 48)), 0), -1)
+            prediction = self.classifier.predict(cropped_img)
+            maxindex = int(np.argmax(prediction))
+            cv2.putText(frame, self.labels[maxindex], (x+20, y-60),
+                        cv2.FONT_HERSHEY_SIMPLEX, 1.4, (0, 0, 255), 2, cv2.LINE_AA)
+
+        return frame
+
+    def run(self):
+        start_time = time.time()
+        if self.is_image:
+            image_arr = cv2.imread(self.input_path)
+            predicted_frame = self.predict(frame=image_arr)
+            cv2.imwrite(self.output_path, predicted_frame)
+        else:
+            cap = cv2.VideoCapture(self.input_path)
+
+            # Get video properties
+            frame_width = int(cap.get(3))
+            frame_height = int(cap.get(4))
+            fps = int(cap.get(5))
+
+            # Define the codec and create a VideoWriter object to save the output video as .MOV
+            fourcc = cv2.VideoWriter_fourcc('a', 'v', 'c', '1')
+            out = cv2.VideoWriter(self.output_path, fourcc, fps,
+                                  (frame_width, frame_height))
+
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+
+                predicted_frame = self.predict(frame=frame)
+                out.write(predicted_frame)
+
+        total_time = str(round(time.time() - start_time, 3)) + "s"
+        return self.output_path, total_time