def __init__(self, init_emb_model=True): self.tf_sess = None self.endpoints = None self.tf_input_image = None people_det_model_prototxt_path = "/home/asoriano/Escritorio/spaceai-evolution/Models/MobileNetSSD_caffe/deploy.prototxt" people_det_model_caffemodel_path = "/home/asoriano/Escritorio/spaceai-evolution/Models/MobileNetSSD_caffe/MobileNetSSD_deploy.caffemodel" people_embeddings_path = "/home/asoriano/Escritorio/spaceai-evolution/Models/person_embedding_tf/checkpoint-25000" # - Reading people detection model self.people_det_model = cv2.dnn.readNetFromCaffe(people_det_model_prototxt_path, people_det_model_caffemodel_path) self.people_det_model_class_names = {0: 'background', 15: 'person'} # - People embeddings model if init_emb_model: tf.Graph().as_default() self.tf_sess = tf.Session() self.tf_input_image = tf.zeros([1, 256, 128, 3], dtype=tf.float32) endpoints, body_prefix = model.endpoints(self.tf_input_image, is_training=False) with tf.name_scope('head'): self.endpoints = head.head(endpoints, 128, is_training=False) tf.train.Saver().restore(self.tf_sess, people_embeddings_path)
import os import cv2 import glob import h5py import tensorflow as tf import numpy as np import nets.resnet_v1_50 as model import heads.fc1024 as head os.environ["CUDA_VISIBLE_DEVICES"] = "3" # Tensorflow descriptor model tf.Graph().as_default() sess = tf.Session() images = tf.zeros([1, 256, 256, 3], dtype=tf.float32) endpoints, body_prefix = model.endpoints(images, is_training=False) with tf.name_scope('head'): endpoints = head.head(endpoints, 128, is_training=False) tf.train.Saver().restore(sess, 'exps/checkpoint-20000') def sim_vector(img): resize_img = cv2.resize(img, (256, 256)) resize_img = np.expand_dims(resize_img, axis=0) emb = sess.run(endpoints['emb'], feed_dict={images: resize_img}) return emb if __name__ == '__main__': dir_images = './oxford/*.jpg'