71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
import cv2
|
|
import dlib
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
from tensorflow.python.keras.models import Model
|
|
from tensorflow.python.keras.applications import ResNet50
|
|
from tensorflow.python.keras.layers import Dense
|
|
|
|
import config
|
|
|
|
|
|
def get_age_model():
|
|
age_model = ResNet50(
|
|
include_top=False,
|
|
weights='imagenet',
|
|
input_shape=(config.RESNET50_DEFAULT_IMG_WIDTH, config.RESNET50_DEFAULT_IMG_WIDTH, 3),
|
|
pooling='avg'
|
|
)
|
|
|
|
prediction = Dense(units=101,
|
|
kernel_initializer='he_normal',
|
|
use_bias=False,
|
|
activation='softmax',
|
|
name='pred_age')(age_model.output)
|
|
|
|
age_model = Model(inputs=age_model.input, outputs=prediction)
|
|
return age_model
|
|
|
|
|
|
def get_model():
|
|
base_model = get_age_model()
|
|
base_model.load_weights(config.AGE_TRAINED_MODEL_PATH)
|
|
print('Loaded weights from age classifier')
|
|
|
|
return base_model
|
|
|
|
def get_trained_model():
|
|
_model = get_model()
|
|
return _model
|
|
|
|
model = get_trained_model()
|
|
#detector = dlib.get_frontal_face_detector()
|
|
|
|
graph = tf.get_default_graph()
|
|
|
|
def predict_age(detected_faces, img):
|
|
input_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
img_h, img_w, _ = np.shape(input_img)
|
|
|
|
faces = np.empty((len(detected_faces), config.RESNET50_DEFAULT_IMG_WIDTH,
|
|
config.RESNET50_DEFAULT_IMG_WIDTH, 3))
|
|
|
|
|
|
for i, d in enumerate(detected_faces):
|
|
x1, y1, x2, y2, w, h = d.left(), d.top(), d.right() + 1, \
|
|
d.bottom() + 1, d.width(), d.height()
|
|
xw1 = max(int(x1 - config.MARGIN * w), 0)
|
|
yw1 = max(int(y1 - config.MARGIN * h), 0)
|
|
xw2 = min(int(x2 + config.MARGIN * w), img_w - 1)
|
|
yw2 = min(int(y2 + config.MARGIN * h), img_h - 1)
|
|
|
|
faces[i, :, :, :] = cv2.resize(img[yw1:yw2 + 1, xw1:xw2 + 1, :], (
|
|
config.RESNET50_DEFAULT_IMG_WIDTH,
|
|
config.RESNET50_DEFAULT_IMG_WIDTH)) / 255.00
|
|
|
|
with graph.as_default():
|
|
results = model.predict(faces)
|
|
ages = np.arange(0, 101).reshape(101, 1)
|
|
predicted_ages = results.dot(ages).flatten()
|
|
return predicted_ages
|