Spaces:
Running
Running
| 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 | |