class Yolow(object):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# sess = tf.Session()
default_models = load_default_models()
model_types = list(default_models.keys())
freeze_dir = 'data/pb/'
tf_models = {model_types[0]: Yolov3, model_types[1]: Yolov3Tiny}
def __init__(self,
model_type,
model_file,
anchor_file,
num_classes,
input_size,
labels,
is_training=False):
if model_type not in self.model_types:
raise ValueError(
'model_type can only be either \'full\' or \'tiny\'.')
elif not model_type:
model_type = self.model_types[0]
self.model_type = model_type
if not model_file:
model_file = './data/bin/{}'.format(
self.default_models.get(model_type))
elif not os.path.exists(model_file):
raise ValueError(
'model file {} does not exist.'.format(model_file))
self.model_file = model_file
if '.pb' not in self.model_file:
self.frozen_filename = '_'.join(
['frozen',
os.path.basename(self.model_file).split('.')[0]])
self.frozen_filename = self.freeze_dir + self.frozen_filename + '.pb'
if not input_size:
input_size = 416
if type(input_size) is int:
self.input_size = input_size, input_size
else:
self.input_size = input_size
self.labels = labels
self.imer = Imager(self.input_size, self.labels)
if os.path.exists(self.frozen_filename):
self.defrost()
self.input = tf.get_default_graph().get_tensor_by_name(
'import/input:0')
self.output = tf.get_default_graph().get_tensor_by_name(
'import/detections/output:0')
else:
if not anchor_file:
anchor_file = 'data/anchors/' + self.model_type + '.txt'
elif not os.path.exists(anchor_file):
raise ValueError(
'{} anchor file does not exist.'.format(anchor_file))
self.anchor_file = anchor_file
self.num_classes = num_classes
self.is_training = is_training
self.input = tf.placeholder(
tf.float32, [None, self.input_size[0], self.input_size[1], 3],
'input')
self.model = self.tf_models[self.model_type](self.input,
self.num_classes,
self.input_size,
self.anchor_file,
self.is_training)
with tf.variable_scope('detections'):
self.output = self.model.graph()
self.loader = WeightLoader(tf.global_variables('detections'),
self.model_file)
# self.sess.run(tf.global_variables_initializer())
self.sess.run(self.loader.load_now())
self.freeze()
def set_input(self, images):
if type(images) == str:
self.imer.imset_from_path(images)
else:
self.imer.imset(images)
def predict(self, confidence_theshold=.6, iou_threshold=.5):
input_list = self.imer.preprocess()
feed_dict = {self.input: input_list}
batch_detections = self.sess.run(self.output, feed_dict)
pred_list = predict(batch_detections, confidence_theshold,
iou_threshold)
return self.imer.visualise_preds(pred_list)
def freeze(self):
graph_def = tf.graph_util.convert_variables_to_constants(
sess=self.sess,
input_graph_def=tf.get_default_graph().as_graph_def(),
output_node_names=['detections/output'])
if not os.path.exists(self.freeze_dir):
os.makedirs(self.freeze_dir)
with tf.gfile.GFile(self.frozen_filename, 'wb') as f:
f.write(graph_def.SerializeToString())
def defrost(self):
print('Found frozen model {}, defrost and use!'.format(
self.frozen_filename))
with tf.gfile.GFile(self.frozen_filename, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def)
class YolowNCS(object):
def __init__(self, model_name_path, input_size, labels, num_requests=2):
self.model = model_name_path + '.xml'
self.weights = model_name_path + '.bin'
self.labels = labels
self.input_size = input_size
self.imer = Imager(self.input_size, self.labels)
if not os.path.exists(self.model) or not os.path.exists(self.weights):
raise ValueError(
'model files {} does not exist.'.format(model_name_path))
self.plugin = IEPlugin(device='MYRIAD')
log.info('Loading network files:\n\t{}\n\t{}'.format(
self.model, self.weights))
self.net = IENetwork(model=self.model, weights=self.weights)
log.info('Preparing inputs')
self.input_blob = next(iter(self.net.inputs))
self.net.batch_size = 1
log.info('Loading model to the plugin')
self.current_request_id = 0
self.next_request_id = 1
self.num_requests = num_requests
self.exec_net = self.plugin.load(network=self.net,
num_requests=self.num_requests)
def set_input(self, images):
if type(images) == str:
self.imer.imset_from_path(images)
else:
self.imer.imset(images)
def predict(self,
confidence_threshold=.5,
iou_threshold=.4,
async_mode=False):
self.out_list = list()
in_frames = self.imer.ncs_preprocess()
for in_frame, orig_frame in zip(in_frames, self.imer.ims):
objects = list()
origin_im_size = orig_frame.shape[:-1]
input_dict = {self.input_blob: in_frame}
request_handle = self.exec_net.requests[self.current_request_id]
if async_mode:
next_request_id = self.current_request_id + 1
if next_request_id == self.num_requests:
next_request_id = 0
else:
next_request_id = self.current_request_id
self.exec_net.start_async(request_id=next_request_id,
inputs=input_dict)
if async_mode:
self.current_request_id = next_request_id
request_handle.wait()
pred_dict = request_handle.outputs
for layer_name, out_blob in pred_dict.items():
params = self.net.layers[layer_name].params
layer_params = YoloV3Params(params, out_blob.shape[2])
objects += parse_yolo_region(out_blob, in_frame.shape[2:],
origin_im_size, layer_params,
confidence_threshold)
for i in range(len(objects)):
if objects[i]['confidence'] == 0:
continue
for j in range(i + 1, len(objects)):
if intersection_over_union(objects[i],
objects[j]) > iou_threshold:
objects[j]['confidence'] = 0
objects = [
obj for obj in objects
if obj['confidence'] >= confidence_threshold
]
self.out_list.append(objects)
out_frames = self.imer.ncs_visualise_preds(self.out_list)
return out_frames
def get_output(self):
try:
return self.out_list
except:
raise ValueError(
'output does not exist, YolowNCS.predict() method must be called prior to this method.'
)
