mirror of
https://github.com/serengil/deepface.git
synced 2025-06-07 20:15:21 +00:00
279 lines
11 KiB
Python
279 lines
11 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Class definition of YOLO_v3 style detection model on image and video
|
|
"""
|
|
|
|
import colorsys
|
|
import os
|
|
import re
|
|
from extension.constants import CONSTANTS
|
|
import numpy as np
|
|
from keras import backend as K
|
|
from keras.models import load_model
|
|
from keras.layers import Input
|
|
from PIL import Image, ImageFont, ImageDraw
|
|
|
|
from yolo3.model import yolo_eval, yolo_body
|
|
from yolo3.utils import letterbox_image
|
|
import os
|
|
from keras.utils import multi_gpu_model
|
|
|
|
# for tensorflow 2.0 emotion model
|
|
import tensorflow as tf
|
|
from tensorflow.keras.applications.xception import preprocess_input
|
|
from keras.preprocessing import image as keras_Image
|
|
|
|
# nms code from deepsort
|
|
from yolo3.preprocessing import non_max_suppression
|
|
|
|
|
|
class YOLO(object):
|
|
emotion_dict = {0: CONSTANTS['EMOTION']['UNSETTLED'], 1: CONSTANTS['EMOTION']['HAPPY'],
|
|
2: CONSTANTS['EMOTION']['NEUTRAL'], 3: CONSTANTS['EMOTION']['SAD'],
|
|
4: CONSTANTS['EMOTION']['ANXIOUS']}
|
|
|
|
_defaults = {
|
|
"model_path": 'yolov3_all/trained_weights_final.h5',
|
|
"anchors_path": 'model_data/yolo3_anchors.txt',
|
|
"classes_path": 'model_data/pet_classes.txt',
|
|
"emotion_model_dir": "emotion/",
|
|
"score": 0.3,
|
|
"iou": 0.45,
|
|
"model_image_size": (608, 608),
|
|
"gpu_num": 1,
|
|
}
|
|
|
|
@classmethod
|
|
def get_defaults(cls, n):
|
|
if n in cls._defaults:
|
|
return cls._defaults[n]
|
|
else:
|
|
return "Unrecognized attribute name '" + n + "'"
|
|
|
|
def __init__(self, **kwargs):
|
|
self.__dict__.update(self._defaults) # set up default values
|
|
self.__dict__.update(kwargs) # and update with user overrides
|
|
self.class_names = self._get_class()
|
|
self.anchors = self._get_anchors()
|
|
self.sess = K.get_session()
|
|
|
|
# init breed model
|
|
self.boxes, self.scores, self.classes = self.generate()
|
|
|
|
# init emotion models
|
|
self.cat_emotion_model, self.dog_emotion_model = self.load_emotion_models()
|
|
|
|
def _get_class(self):
|
|
classes_path = os.path.expanduser(self.classes_path)
|
|
with open(classes_path) as f:
|
|
class_names = f.readlines()
|
|
class_names = [c.strip() for c in class_names]
|
|
return class_names
|
|
|
|
def _get_anchors(self):
|
|
anchors_path = os.path.expanduser(self.anchors_path)
|
|
with open(anchors_path) as f:
|
|
anchors = f.readline()
|
|
anchors = [float(x) for x in anchors.split(',')]
|
|
return np.array(anchors).reshape(-1, 2)
|
|
|
|
def load_emotion_models(self):
|
|
cat_model_path = os.path.join(self.emotion_model_dir, "Cat_classifier.h5")
|
|
dog_model_path = os.path.join(self.emotion_model_dir, "Dog_classifier.h5")
|
|
|
|
cat = tf.keras.models.load_model(cat_model_path)
|
|
self.graph_cat = tf.get_default_graph()
|
|
|
|
dog = tf.keras.models.load_model(dog_model_path)
|
|
self.graph_dog = tf.get_default_graph()
|
|
|
|
cat.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
|
|
dog.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
|
|
|
|
return cat, dog
|
|
|
|
def generate(self):
|
|
model_path = os.path.expanduser(self.model_path)
|
|
assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'
|
|
|
|
# Load model, or construct model and load weights.
|
|
num_anchors = len(self.anchors)
|
|
num_classes = len(self.class_names)
|
|
try:
|
|
self.yolo_model = load_model(model_path, compile=False)
|
|
except:
|
|
self.yolo_model = yolo_body(Input(shape=(None, None, 3)), num_anchors // 3, num_classes)
|
|
self.yolo_model.load_weights(self.model_path) # make sure model, anchors and classes match
|
|
else:
|
|
assert self.yolo_model.layers[-1].output_shape[-1] == \
|
|
num_anchors / len(self.yolo_model.output) * (num_classes + 5), \
|
|
'Mismatch between model and given anchor and class sizes'
|
|
|
|
print('{} model, anchors, and classes loaded.'.format(model_path))
|
|
|
|
# Generate colors for drawing bounding boxes.
|
|
hsv_tuples = [(x / len(self.class_names), 1., 1.)
|
|
for x in range(len(self.class_names))]
|
|
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
|
|
self.colors = list(
|
|
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
|
|
self.colors))
|
|
np.random.seed(10101) # Fixed seed for consistent colors across runs.
|
|
np.random.shuffle(self.colors) # Shuffle colors to decorrelate adjacent classes.
|
|
np.random.seed(None) # Reset seed to default.
|
|
|
|
# Generate output tensor targets for filtered bounding boxes.
|
|
self.input_image_shape = K.placeholder(shape=(2,))
|
|
if self.gpu_num >= 2:
|
|
self.yolo_model = multi_gpu_model(self.yolo_model, gpus=self.gpu_num)
|
|
boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
|
|
len(self.class_names), self.input_image_shape,
|
|
score_threshold=self.score, iou_threshold=self.iou)
|
|
return boxes, scores, classes
|
|
|
|
def detect_image(self, image):
|
|
|
|
if self.model_image_size != (None, None):
|
|
assert self.model_image_size[0] % 32 == 0, 'Multiples of 32 required'
|
|
assert self.model_image_size[1] % 32 == 0, 'Multiples of 32 required'
|
|
boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
|
|
else:
|
|
new_image_size = (image.width - (image.width % 32),
|
|
image.height - (image.height % 32))
|
|
boxed_image = letterbox_image(image, new_image_size)
|
|
image_data = np.array(boxed_image, dtype='float32')
|
|
|
|
image_data /= 255.
|
|
image_data = np.expand_dims(image_data, 0) # Add batch dimension.
|
|
|
|
out_boxes, out_scores, out_classes = self.sess.run(
|
|
[self.boxes, self.scores, self.classes],
|
|
feed_dict={
|
|
self.yolo_model.input: image_data,
|
|
self.input_image_shape: [image.size[1], image.size[0]]
|
|
# K.learning_phase(): 0
|
|
})
|
|
|
|
# non-maxima suppression ACROSS classes
|
|
indices = non_max_suppression(np.array(out_boxes),
|
|
self.iou,
|
|
out_scores)
|
|
|
|
out_boxes = [out_boxes[i] for i in indices]
|
|
out_scores = [out_scores[i] for i in indices]
|
|
out_classes = [out_classes[i] for i in indices]
|
|
|
|
font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
|
|
size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
|
|
thickness = (image.size[0] + image.size[1]) // 300
|
|
|
|
detection_dicts = []
|
|
|
|
count = 0
|
|
|
|
showid = len(out_classes) > 1
|
|
|
|
for i, c in reversed(list(enumerate(out_classes))):
|
|
|
|
# prediction of breed and bounding box
|
|
predicted_class = self.class_names[c]
|
|
box = out_boxes[i]
|
|
score = out_scores[i]
|
|
|
|
# bounding box corners
|
|
top, left, bottom, right = box
|
|
top = max(0, np.floor(top + 0.5).astype('int32'))
|
|
left = max(0, np.floor(left + 0.5).astype('int32'))
|
|
bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
|
|
right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
|
|
|
|
# crop image head and pass to emotion model
|
|
head_image = image.crop([left, top, right, bottom])
|
|
head_image = head_image.resize((200, 200))
|
|
head_image = keras_Image.img_to_array(head_image)
|
|
head_image = np.expand_dims(head_image, axis=0)
|
|
head_image = preprocess_input(head_image)
|
|
|
|
if predicted_class[0].islower():
|
|
# lowercase: dog
|
|
with self.graph_dog.as_default():
|
|
prediction = self.dog_emotion_model.predict(head_image)[0]
|
|
else:
|
|
# uppercase: cat
|
|
with self.graph_cat.as_default():
|
|
prediction = self.cat_emotion_model.predict(head_image)[0]
|
|
|
|
prediction = np.round(prediction, decimals=2)
|
|
# prediction = [float(np.round(p * 100, 2)) for p in prediction]
|
|
prediction_ = np.argmax(prediction)
|
|
|
|
# label
|
|
label = 'ID: {}'.format(count)
|
|
draw = ImageDraw.Draw(image)
|
|
label_size = draw.textsize(label, font=font)
|
|
|
|
if top - label_size[1] >= 0:
|
|
text_origin = np.array([left, top - label_size[1]])
|
|
else:
|
|
text_origin = np.array([left, top + 1])
|
|
|
|
# My kingdom for a good redistributable image drawing library.
|
|
for i in range(thickness):
|
|
draw.rectangle(
|
|
[left + i, top + i, right - i, bottom - i],
|
|
outline=self.colors[c])
|
|
|
|
if showid:
|
|
draw.rectangle(
|
|
[tuple(text_origin), tuple(text_origin + label_size)],
|
|
fill=self.colors[c]
|
|
)
|
|
draw.text(text_origin, label, fill=(0, 0, 0), font=font)
|
|
del draw
|
|
|
|
# put results into dict list
|
|
if predicted_class[0].isupper():
|
|
pet = "Cat"
|
|
else:
|
|
pet = "Dog"
|
|
|
|
all_emotions = [
|
|
{'name': CONSTANTS['EMOTION']['UNSETTLED'], 'value': prediction[0]},
|
|
{'name': CONSTANTS['EMOTION']['HAPPY'], 'value': prediction[1]},
|
|
{'name': CONSTANTS['EMOTION']['NEUTRAL'], 'value': prediction[2]},
|
|
{'name': CONSTANTS['EMOTION']['SAD'], 'value': prediction[3]},
|
|
{'name': CONSTANTS['EMOTION']['ANXIOUS'], 'value': prediction[4]},
|
|
]
|
|
all_emotions.sort(key=lambda x: x['value'], reverse=True)
|
|
all_emotions_result = []
|
|
for each_emotion in all_emotions:
|
|
all_emotions_result.append(
|
|
each_emotion['name'] + ' (' + str(np.round(each_emotion['value'] * 100)) + '%)')
|
|
|
|
predicted_breed = re.sub(r"_", " ", predicted_class, flags=re.IGNORECASE)
|
|
obj_dict = {
|
|
'id': count,
|
|
'pet': pet,
|
|
'breed': "{} ({}%)".format(predicted_breed.title(), np.round(score * 100)),
|
|
'breedTitle': predicted_breed.title(),
|
|
'breedScore': np.round(score * 100),
|
|
'emotion': {
|
|
'mostLikelyTitle': self.emotion_dict[prediction_],
|
|
'mostLikelyScore': np.round(np.max(prediction) * 100),
|
|
'mostLikely': "{} ({}%)".format(self.emotion_dict[prediction_], np.round(np.max(prediction) * 100)),
|
|
'allEmotions': all_emotions_result
|
|
},
|
|
}
|
|
|
|
count += 1
|
|
|
|
detection_dicts.append(obj_dict)
|
|
|
|
# predicted_img_name = "predicted_image.jpg"
|
|
# image.save(predicted_img_name, "JPEG")
|
|
|
|
return image, detection_dicts
|
|
|
|
def close_session(self):
|
|
self.sess.close()
|