def __init__(self, path=DEFAULT_MODEL_PATH): logger.info('Loading model from: {}...'.format(path)) # for face detection self.detector = MTCNN() try: os.mkdir('./img') except OSError: pass # load model and weights self.img_size = 64 self.stage_num = [3, 3, 3] self.lambda_local = 1 self.lambda_d = 1 # load pre-trained model self.model = SSR_net(self.img_size, self.stage_num, self.lambda_local, self.lambda_d)() self.model.load_weights(path) self.graph = tf.get_default_graph() logger.info('Loaded model')
# import sys # from PIL import Image # import pygame # from keras import backend # import tensorflow as tf # import matplotlib # from matplotlib import pyplot as plt # import matplotlib.patches as patches # from moviepy.editor import * # weight path weight_file = "./assets/ssrnet_3_3_3_64_1.0_1.0.h5" model = SSR_net(64, [3, 3, 3], 1, 1)() model.load_weights(weight_file) detector = MTCNN() def draw_label(image, point, label, font=cv2.FONT_HERSHEY_SIMPLEX, font_scale=1, thickness=2): size = cv2.getTextSize(label, font, font_scale, thickness)[0] x, y = point cv2.rectangle(image, (x, y - size[1]), (x + size[0], y), (255, 0, 0), cv2.FILLED)
class ModelWrapper(MAXModelWrapper): MODEL_META_DATA = { 'id': 'ssrnet', 'name': 'SSR-Net Facial Age Estimator Model', 'description': 'SSR-Net Facial Recognition and Age Prediction model; trained using Keras on the IMDB-WIKI dataset', 'type': 'Facial Recognition', 'source': 'https://developer.ibm.com/exchanges/models/all/max-facial-age-estimator/', 'license': 'MIT' } def __init__(self, path=DEFAULT_MODEL_PATH): logger.info('Loading model from: {}...'.format(path)) # for face detection self.detector = MTCNN() try: os.mkdir('./img') except OSError: pass # load model and weights self.img_size = 64 self.stage_num = [3, 3, 3] self.lambda_local = 1 self.lambda_d = 1 # load pre-trained model self.model = SSR_net(self.img_size, self.stage_num, self.lambda_local, self.lambda_d)() self.model.load_weights(path) self.graph = tf.get_default_graph() logger.info('Loaded model') def _pre_process(self, input_img): ad = 0.4 img_h, img_w, _ = np.shape(input_img) # check image w/h > 1024 input_img = img_resize(input_img) img_h, img_w, _ = np.shape(input_img) detected = self.detector.detect_faces(input_img) faces = np.empty((len(detected), self.img_size, self.img_size, 3)) for i, d in enumerate(detected): if d['confidence'] > 0.85: x1, y1, w, h = d['box'] x2 = x1 + w y2 = y1 + h # Normalized bbx coordinates d['box'] = [float(y1) / img_h, float(x1) / img_w, float(y2) / img_h, float(x2) / img_w] xw1 = max(int(x1 - ad * w), 0) yw1 = max(int(y1 - ad * h), 0) xw2 = min(int(x2 + ad * w), img_w - 1) yw2 = min(int(y2 + ad * h), img_h - 1) faces[i, :, :, :] = cv2.resize(input_img[yw1:yw2 + 1, xw1:xw2 + 1, :], (self.img_size, self.img_size)) return (faces, detected) def _predict(self, pre_x): faces=pre_x[0] with self.graph.as_default(): predicted_ages = self.model.predict(faces) return (predicted_ages,pre_x[1]) def _post_process(self,post_rst): predicted_ages=post_rst[0] detected=post_rst[1] pred_res = [] for i, d in enumerate(detected): if d['confidence'] > 0.85: pre_age=predicted_ages[i].astype(int) pred_res.append([{'box': d['box'], 'age':pre_age}]) return pred_res
class ModelWrapper(object): """Model wrapper for SavedModel format""" def __init__(self, path=DEFAULT_MODEL_PATH): logger.info('Loading model from: {}...'.format(path)) # for face detection self.detector = MTCNN() try: os.mkdir('./img') except OSError: pass # load model and weights self.img_size = 64 self.stage_num = [3, 3, 3] self.lambda_local = 1 self.lambda_d = 1 # load pre-trained model self.model = SSR_net(self.img_size, self.stage_num, self.lambda_local, self.lambda_d)() self.model.load_weights(path) self.graph = tf.get_default_graph() logger.info('Loaded model') def predict(self, x): input_img = x # python version pyFlag = '' if len(sys.argv) < 3: pyFlag = '2' # default to use moviepy to show, this can work on python2.7 and python3.5 elif len(sys.argv) == 3: pyFlag = sys.argv[2] # python version else: print('Wrong input!') sys.exit() detected = '' # make this not local variable ad = 0.4 if pyFlag == '3': input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB) img_h, img_w, _ = np.shape(input_img) input_img = cv2.resize(input_img, (1024, int(1024 * img_h / img_w))) img_h, img_w, _ = np.shape(input_img) detected = self.detector.detect_faces(input_img) faces = np.empty((len(detected), self.img_size, self.img_size, 3)) for i, d in enumerate(detected): if d['confidence'] > 0.95: x1, y1, w, h = d['box'] x2 = x1 + w y2 = y1 + h xw1 = max(int(x1 - ad * w), 0) yw1 = max(int(y1 - ad * h), 0) xw2 = min(int(x2 + ad * w), img_w - 1) yw2 = min(int(y2 + ad * h), img_h - 1) cv2.rectangle(input_img, (x1, y1), (x2, y2), (255, 0, 0), 2) # cv2.rectangle(img, (xw1, yw1), (xw2, yw2), (255, 0, 0), 2) faces[i, :, :, :] = cv2.resize(input_img[yw1:yw2 + 1, xw1:xw2 + 1, :], (self.img_size, self.img_size)) if len(detected) > 0: with self.graph.as_default(): predicted_ages = self.model.predict(faces) # prediction results with BBX & AGES pred_res = [] for i, d in enumerate(detected): if d['confidence'] > 0.8: pre_age=predicted_ages[i].astype(int) pred_res.append([{'box': d['box'], 'age':pre_age}]) return pred_res