class ResNet50ExtractorAPI:
def __init__(self, graph_prefix='', flush_pipe_on_read=False):
self.__flush_pipe_on_read = flush_pipe_on_read
self.__feature_dim = (2048)
self.__image_shape = (224, 224, 3)
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
if not graph_prefix:
self.__graph_prefix = ''
else:
self.__graph_prefix = graph_prefix + '/'
def __preprocess(self, original):
preprocessed = cv2.resize(original,
tuple(self.__image_shape[:2][::-1]))
preprocessed = resnet50.preprocess_input(preprocessed)
return preprocessed
def __in_pipe_process(self, inference):
images = inference.get_input()
if type(images) is not list:
images = [images]
data = []
for img in images:
data.append(self.__preprocess(img))
data = np.array(data)
inference.set_data(data)
return inference
def __out_pipe_process(self, result):
result, inference = result
inference.set_result(result)
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
K.set_session(session_runner.get_session())
self.__tf_sess = K.get_session()
with self.__tf_sess.as_default():
self.__model = resnet50.ResNet50(weights='imagenet',
include_top=False,
pooling='avg')
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__session_runner.get_in_pipe().push(
(self.__job, inference))
else:
self.__in_pipe.wait()
def __job(self, inference):
try:
self.__out_pipe.push(
(self.__model.predict(inference.get_data()), inference))
except:
self.__out_pipe.push((np.zeros((0, self.__feature_dim),
np.float32), inference))
class Sort:
def __init__(self, max_age=120, min_hits=3, flush_pipe_on_read=False):
"""
Sets key parameters for SORT
"""
self.max_age = max_age
self.min_hits = min_hits
self.trackers = []
self.frame_count = 0
self.__flush_pipe_on_read = flush_pipe_on_read
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__outpass_pipe = Pipe(self.__out_pipe_process)
def __in_pipe_process(self, inference):
detections = inference.get_boxes_tlbr(normalized=False)
frame = inference.get_input()
classes = inference.get_classes()
person_detections = []
scores = inference.get_scores()
for i in range(len(classes)):
if classes[i] == inference.get_category(
'person') and scores[i] > .5:
person_detections.append([
detections[i][1], detections[i][0], detections[i][3],
detections[i][2]
])
return (frame, person_detections)
def __out_pipe_process(self, inference):
image, f_vecs, bboxes = inference
if bboxes:
matched, unmatched_dets, unmatched_trks = Tracker.associate_detections_to_trackers(
f_vecs, self.trackers)
# update matched trackers with assigned detections
for t, trk in enumerate(self.trackers):
if (t not in unmatched_trks):
d = matched[np.where(matched[:, 1] == t)[0], 0][0]
trk.update(
bboxes[d],
f_vecs[d]) ## for dlib re-intialize the trackers ?!
# create and initialise new trackers for unmatched detections
for i in unmatched_dets:
trk = Tracker(bboxes[i], f_vecs[i])
self.trackers.append(trk)
i = len(self.trackers)
ret = []
for trk in reversed(self.trackers):
# if bboxes == []:
# trk.update([], img)
d = trk.get_bbox()
if (trk.get_hit_streak() >= self.min_hits
or self.frame_count <= self.min_hits):
ret.append(np.concatenate((d, [trk.get_id()])).reshape(
1, -1)) # +1 as MOT benchmark requires positive
i -= 1
# remove dead tracklet
if (trk.get_time_since_update() > self.max_age):
self.trackers.pop(i)
print(len(self.trackers))
if (len(ret) > 0):
return np.concatenate(ret), image
return np.empty((0, 5)), image
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_session_runner(self, model, session_runner):
self.__session_runner = session_runner
# K.set_session(session_runner.get_session())
# self.__tf_sess = K.get_session()
self.__image_encoder = ImageEncoder(model, session_runner)
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, data_tuple = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.data_tuple = data_tuple
self.__session_runner.add_job(self.__job())
else:
self.__in_pipe.pull_wait()
def __job(self):
f_vecs = self.__image_encoder.extract_features(self.data_tuple[0],
self.data_tuple[1])
self.__out_pipe.push((self.data_tuple[0], f_vecs, self.data_tuple[1]))
def update(self, image, bboxes):
"""
class MarsExtractorAPI:
def __init__(self, graph_prefix='', flush_pipe_on_read=False):
self.__flush_pipe_on_read = flush_pipe_on_read
self.__model_path = pretrained_path.get() + "/mars-small128.pb"
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
if not graph_prefix:
self.__graph_prefix = ''
else:
self.__graph_prefix = graph_prefix + '/'
def get_input_shape(self):
return self.__image_shape
def __preprocess(self, original):
preprocessed = cv2.resize(original,
tuple(self.__image_shape[:2][::-1]))
# preprocessed = resnet50.preprocess_input(preprocessed)
return preprocessed
def __in_pipe_process(self, inference):
images = inference.get_input()
if type(images) is not list:
images = [images]
data = []
for img in images:
data.append(self.__preprocess(img))
# data = np.array(data)
inference.set_data(data)
return inference
def __out_pipe_process(self, result):
result, inference = result
inference.set_result(result)
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
self.__tf_sess = session_runner.get_session()
with self.__tf_sess.graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.__model_path, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name=self.__graph_prefix)
self.__input_var = tf.get_default_graph().get_tensor_by_name(
self.__graph_prefix + "images:0")
self.__output_var = tf.get_default_graph().get_tensor_by_name(
self.__graph_prefix + "features:0")
assert len(self.__output_var.get_shape()) == 2
assert len(self.__input_var.get_shape()) == 4
self.__feature_dim = self.__output_var.get_shape().as_list()[-1]
self.__image_shape = self.__input_var.get_shape().as_list()[1:]
# print(self.__image_shape)
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__session_runner.get_in_pipe().push(
SessionRunnable(self.__job, inference))
else:
self.__in_pipe.pull_wait()
def __job(self, inference):
x = inference.get_data()
if len(x) > 0:
self.__out_pipe.push(
(self.__tf_sess.run(self.__output_var,
feed_dict={self.__input_var:
x}), inference))
else:
self.__out_pipe.push((np.zeros((0, self.__feature_dim),
np.float32), inference))
class AgeDetection(object):
'''
Path to the trained model path for age_detection
'''
WRN_WEIGHTS_PATH = age_model_path.get() + '/age_model.h5'
def __init__(self, flush_pipe_on_read=False):
self.__detector = None
self.__face_size = 64
self.__flush_pipe_on_read = flush_pipe_on_read
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
def crop_face(self, imgarray, section, margin=40, size=64):
"""
:param imgarray: full image
:param section: face detected area (x, y, w, h)
:param margin: add some margin to the face detected area to include a full head
:param size: the result image resolution with be (size x size)
:return: resized image in numpy array with shape (size x size x 3)
"""
img_h, img_w, _ = imgarray.shape
if section is None:
section = [0, 0, img_w, img_h]
(x, y, w, h) = section
margin = int(min(w, h) * margin / 100)
x_a = x - margin
y_a = y - margin
x_b = x + w + margin
y_b = y + h + margin
if x_a < 0:
x_b = min(x_b - x_a, img_w - 1)
x_a = 0
if y_a < 0:
y_b = min(y_b - y_a, img_h - 1)
y_a = 0
if x_b > img_w:
x_a = max(x_a - (x_b - img_w), 0)
x_b = img_w
if y_b > img_h:
y_a = max(y_a - (y_b - img_h), 0)
y_b = img_h
cropped = imgarray[y_a:y_b, x_a:x_b]
resized_img = cv2.resize(cropped, (size, size),
interpolation=cv2.INTER_AREA)
resized_img = np.array(resized_img)
return resized_img, (x_a, y_a, x_b - x_a, y_b - y_a)
def __in_pipe_process(self, inference):
'''
:param inference: inference object with input frame
:return: inference: inference object with age_inference object containing face bbox details,
data condaining np_array of all faces detected
'''
frame = inference.get_input()
faces = self.__detector.detect_faces(frame)
faces = list_filter(faces, 0.90)
# print("faces " ,faces)
bboxes = []
face_imgs = np.empty(
(len(faces), self.__face_size, self.__face_size, 3))
for i, face in enumerate(faces):
face_img, cropped = self.crop_face(frame,
face[:4],
margin=40,
size=self.__face_size)
(x, y, w, h) = cropped
bboxes.append([x, y, x + w, y + h])
# face_img = (face_img/255)-0.5
# cv2.rectangle(self.image_bounding_boxes, (x, y), (x + w, y + h), (255, 200, 0), 2)
face_imgs[i, :, :, :] = face_img
age_inference = AgeInference(frame, bboxes=bboxes)
inference.set_result(age_inference)
inference.set_data(face_imgs)
return inference
def __out_pipe_process(self, inference):
'''
:param inference: inference object with faces np array in data,
age_inference object with bboxes in result
:return: inference: inference object with age_inference object in result after processing
face np array from data for age, gender and ethnicity
'''
# placeholder for cropped faces
results, inference = inference
age_inference = inference.get_result()
if results is not None:
predicted_genders = results[0]
ages = np.arange(0, 101).reshape(101, 1)
predicted_ages = results[1].dot(ages).flatten()
predicted_ethnicity = np.argmax(results[2], axis=1)
age_inference.set_ages(predicted_ages)
age_inference.set_genders(predicted_genders)
age_inference.set_ethnicity(predicted_ethnicity)
inference.set_result(age_inference)
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_threading(self, run_on_thread=True):
self.__run_session_on_thread = run_on_thread
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
K.set_session(self.__session_runner.get_session())
self.__tf_sess = K.get_session()
self.__detector = MTCNN(session_runner=self.__session_runner)
with self.__tf_sess.as_default():
self.__model = WideResNet(self.__face_size)()
self.__model.load_weights(AgeDetection.WRN_WEIGHTS_PATH)
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__session_runner.get_in_pipe().push(
SessionRunnable(self.__job, inference))
else:
self.__in_pipe.pull_wait()
def __job(self, inference):
'''
:param inference: run the model on data from inference object and push it to out_pipe
:return:
'''
try:
self.__out_pipe.push(
(self.__model.predict(inference.get_data()), inference))
except:
self.__out_pipe.push((None, inference))
class ImageEncoder:
@staticmethod
def __fetch_model_path(model_name):
return "../resources/networks/" + model_name
def __init__(self,
session_runner,
model_name=PRETRAINED_mars_small128,
graph_prefix=None):
self.__path_to_frozen_graph = self.__fetch_model_path(model_name)
if not graph_prefix:
self.__graph_prefix = ''
else:
self.__graph_prefix = graph_prefix + '/'
self.__session_runner = session_runner
self.__tf_sess = session_runner.get_session()
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
self.init_graph()
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def __in_pipe_process(self, obj):
return obj
def __out_pipe_process(self, obj):
return obj
def init_graph(self):
with self.__tf_sess.graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.__path_to_frozen_graph, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name=self.__graph_prefix)
self.input_var = self.__tf_sess.graph.get_tensor_by_name(
self.__graph_prefix + 'images:0')
self.output_var = self.__tf_sess.graph.get_tensor_by_name(
self.__graph_prefix + 'features:0')
assert len(self.output_var.get_shape()) == 2
assert len(self.input_var.get_shape()) == 4
self.feature_dim = self.output_var.get_shape().as_list()[-1]
self.image_shape = self.input_var.get_shape().as_list()[1:]
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, data_x = self.__in_pipe.pull()
if ret:
self.data_x = data_x
self.__session_runner.add_job(self.__job())
else:
self.__in_pipe.pull_wait()
def __job(self):
__tracking_features = np.zeros((len(self.data_x), self.feature_dim),
np.float32)
self._run_in_batches(
lambda x: self.__tf_sess.run(self.output_var, feed_dict=x),
{self.input_var: self.data_x},
__tracking_features,
batch_size=32)
self.__out_pipe.push(__tracking_features)
def _run_in_batches(self, f, data_dict, out, batch_size):
data_len = len(out)
num_batches = int(data_len / batch_size)
s, e = 0, 0
for i in range(num_batches):
s, e = i * batch_size, (i + 1) * batch_size
batch_data_dict = {k: v[s:e] for k, v in data_dict.items()}
out[s:e] = f(batch_data_dict)
if e < len(out):
batch_data_dict = {k: v[e:] for k, v in data_dict.items()}
out[e:] = f(batch_data_dict)
class YOLOObjectDetectionAPI():
_defaults = {
"model_path": pretrained_path.get() + '/yolo_v3.h5',
"anchors_path": data_path.get() + '/yolo_anchors.txt',
"classes_path": data_path.get() + '/coco.names',
"score": 0.3,
"iou": 0.45,
"model_image_size": (416, 416),
"gpu_num": 1,
}
@staticmethod
def __get_dir_path():
return dirname(realpath(__file__))
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
def __init__(self, graph_prefix=None, flush_pipe_on_read=False):
self.__dict__.update(self._defaults) # set up default values
# self.__dict__.update(kwargs) # and update with user overrides
YOLOObjectDetectionAPI.class_names, self.__category_dict = self._get_class(
)
YOLOObjectDetectionAPI.anchors = self._get_anchors()
self.__graph_prefix = graph_prefix
self.__flush_pipe_on_read = flush_pipe_on_read
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
self.__run_session_on_thread = False
def use_threading(self, run_on_thread=True):
self.__run_session_on_thread = run_on_thread
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
K.set_session(session_runner.get_session())
self.__tf_sess = K.get_session()
self.boxes, self.scores, self.classes = self.generate()
def __in_pipe_process(self, inference):
image = inference.get_input()
pil_image = Image.fromarray(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(
pil_image, tuple(reversed(self.model_image_size)))
else:
new_image_size = (pil_image.width - (pil_image.width % 32),
pil_image.height - (pil_image.height % 32))
boxed_image = letterbox_image(pil_image, new_image_size)
data = np.array(boxed_image, dtype='float32')
data /= 255.
data = np.expand_dims(data, 0) # Add batch dimension.
inference.set_data(data)
inference.get_meta_dict()['PIL'] = pil_image
return inference
# return image
def __out_pipe_process(self, result):
(out_boxes, out_classes, out_scores), inference = result
out_boxes = np.array([[0 if y < 0 else y for y in x]
for x in out_boxes])
result = InferedDetections(inference.get_input(),
len(out_boxes),
out_boxes,
out_classes,
out_scores,
masks=None,
is_normalized=False,
get_category_fnc=self.get_category,
annotator=self.annotate)
inference.set_result(result)
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def _get_class(self):
dir_path = YOLOObjectDetectionAPI.__get_dir_path()
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]
category_dict = {}
for id, name in enumerate(class_names):
category_dict[id] = name
category_dict[name] = id
return class_names, category_dict
def _get_anchors(self):
dir_path = YOLOObjectDetectionAPI.__get_dir_path()
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 generate(self):
dir_path = YOLOObjectDetectionAPI.__get_dir_path()
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)
is_tiny_version = num_anchors == 6 # default setting
try:
self.yolo_model = load_model(model_path, compile=False)
except:
self.yolo_model = tiny_yolo_body(Input(shape=(None, None, 3)), num_anchors // 2, num_classes) \
if is_tiny_version else 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))]
YOLOObjectDetectionAPI.colors = list(
map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
YOLOObjectDetectionAPI.colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
YOLOObjectDetectionAPI.colors))
np.random.seed(10101) # Fixed seed for consistent colors across runs.
np.random.shuffle(YOLOObjectDetectionAPI.colors
) # Shuffle colors to decorrelate adjacent classes.
np.random.shuffle(YOLOObjectDetectionAPI.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 freeze_session(self,
session,
keep_var_names=None,
output_names=None,
clear_devices=True):
"""
Freezes the state of a session into a pruned computation graph.
Creates a new computation graph where variable nodes are replaced by
constants taking their current value in the session. The new graph will be
pruned so subgraphs that are not necessary to compute the requested
outputs are removed.
@param session The TensorFlowsion(detected_boxes, confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
draw_boxes(filtered_boxes, img, classes, (FLAGS.size, FLAGS.size))
img.save(FLAGS.output_img) session to be frozen.
@param keep_var_names A list of variable names that should not be frozen,
or None to freeze all the variables in the graph.
@param output_names Names of the relevant graph outputs.
@param clear_devices Remove the device directives from the graph for better portability.
@return The frozen graph definition.
"""
from tensorflow.python.framework.graph_util import convert_variables_to_constants
graph = session.graph
with graph.as_default():
freeze_var_names = list(
set(v.op.name
for v in tf.global_variables()).difference(keep_var_names
or []))
output_names = output_names or []
output_names += [v.op.name for v in tf.global_variables()]
input_graph_def = graph.as_graph_def()
if clear_devices:
for node in input_graph_def.node:
node.device = ""
frozen_graph = convert_variables_to_constants(
session, input_graph_def, output_names, freeze_var_names)
return frozen_graph
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__session_runner.get_in_pipe().push(
SessionRunnable(
self.__job,
inference,
run_on_thread=self.__run_session_on_thread))
else:
self.__in_pipe.pull_wait()
def __job(self, inference):
data = inference.get_data()
pil_image = inference.get_meta_dict()['PIL']
# image = inference.get_input()
out_boxes, out_scores, out_classes = self.__tf_sess.run(
[self.boxes, self.scores, self.classes],
feed_dict={
self.yolo_model.input: data,
self.input_image_shape: [pil_image.size[1], pil_image.size[0]],
})
self.__out_pipe.push(((out_boxes, out_classes, out_scores), inference))
def close_session(self):
self.__tf_sess.close()
@staticmethod
def annotate(inference):
annotated = inference.image.copy()
image = Image.fromarray(inference.get_image())
font = ImageFont.truetype(font='arial.ttf',
size=np.floor(3e-2 * image.size[1] +
0.5).astype('int32'))
thickness = 1 #(image.size[0] + image.size[1]) // 300
for i, c in reversed(list(enumerate(inference.get_classes()))):
predicted_class = YOLOObjectDetectionAPI.class_names[c]
box = inference.get_boxes_tlbr(normalized=False)[i]
score = inference.get_scores()[i]
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
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'))
# print(label, (left, top), (right, bottom))
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=YOLOObjectDetectionAPI.colors[c])
draw.rectangle(
[tuple(text_origin),
tuple(text_origin + label_size)],
fill=YOLOObjectDetectionAPI.colors[c])
draw.text(text_origin, label, fill=(0, 0, 0), font=font)
del draw
return np.array(image)
def get_category(self, category):
return self.__category_dict[category]
class OFISTObjectTrackingAPI:
def __init__(self, max_age=10000, min_hits=5, flush_pipe_on_read=False):
self.max_age = max_age
self.min_hits = min_hits
self.trackers = []
self.frame_count = 0
self.__bg_frame = None
self.__bg_gray = None
self.__flush_pipe_on_read = flush_pipe_on_read
self.__feature_dim = (2048)
self.__image_shape = (224, 224, 3)
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
number = 0
def __extract_image_patch(self, image, bbox, patch_shape):
sx, sy, ex, ey = np.array(bbox).astype(np.int)
# dx = ex-sx
# dx = int(.25*dx)
# dy = ey-sy
# dy = int(.6*dy)
dx = 0
dy = 0
image = image[sy:ey - dy, sx + dx:ex - dx]
image = cv2.resize(image, tuple(patch_shape[::-1]))
# img_yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
# img_yuv[:, :, 0] = cv2.equalizeHist(img_yuv[:, :, 0])
# image = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
image[0] = cv2.equalizeHist(image[0])
image[1] = cv2.equalizeHist(image[1])
image[2] = cv2.equalizeHist(image[2])
return image
def __in_pipe_process(self, inference):
i_dets = inference.get_input()
frame = i_dets.get_image()
classes = i_dets.get_classes()
boxes = i_dets.get_boxes_tlbr(normalized=False)
masks = i_dets.get_masks()
bboxes = []
scores = i_dets.get_scores()
for i in range(len(classes)):
if classes[i] == i_dets.get_category('person') and scores[i] > .75:
bboxes.append(
[boxes[i][1], boxes[i][0], boxes[i][3], boxes[i][2]])
patches = []
# if self.__bg_frame is None:
# self.__bg_frame = frame
# self.__bg_gray = cv2.cvtColor(self.__bg_frame, cv2.COLOR_BGR2GRAY)
# self.__bg_gray = cv2.GaussianBlur(self.__bg_gray, (5, 5), 0)
#
# gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# gray_frame = cv2.GaussianBlur(gray_frame, (5, 5), 0)
# difference = cv2.absdiff(self.__bg_gray, gray_frame)
# ret, mask = cv2.threshold(difference, 50, 255, cv2.THRESH_BINARY)
# mask = np.stack((mask, mask, mask), axis=2)
#
# image = np.multiply(frame, mask)
#
# inference.get_meta_dict()['mask'] = mask
# inference.get_meta_dict()['diff_img']=image
image = frame
# blur = cv2.GaussianBlur(image, (5, 5), 0)
# image = cv2.addWeighted(image, 1.5, image, -0.5, 0)
# image = cv2.cvtColor(frame, cv2.COLOR_BGR2HLS)
for i in range(len(bboxes)):
box = bboxes[i]
# mask = masks[i]
# mask = np.stack((mask, mask, mask), axis=2)
# image = np.multiply(frame, mask)
patch = self.__extract_image_patch(image, box,
self.__image_shape[:2])
if patch is None:
print("WARNING: Failed to extract image patch: %s." % str(box))
patch = np.random.uniform(0., 255.,
self.__image_shape).astype(np.uint8)
patches.append(patch)
inference.set_data(patches)
inference.get_meta_dict()['bboxes'] = bboxes
return inference
def __out_pipe_process(self, inference):
f_vecs = inference.get_result()
# print(f_vecs.shape)
inference = inference.get_meta_dict()['inference']
bboxes = inference.get_meta_dict()['bboxes']
self.frame_count = min(self.frame_count + 1, 1000)
matched, unmatched_dets, unmatched_trks = KNNTracker.associate_detections_to_trackers(
f_vecs, self.trackers, bboxes, distance_threshold=0.575)
if bboxes:
# # update matched trackers with assigned detections
# for t, trk in enumerate(self.trackers):
# if (t not in unmatched_trks):
# d = matched[np.where(matched[:, 1] == t)[0], 0][0]
# trk.update(bboxes[d], f_vecs[d]) ## for dlib re-intialize the trackers ?!
# update matched trackers with assigned detections
for trk in self.trackers:
if (trk.get_id() not in unmatched_trks):
d = matched[np.where(matched[:, 1] == trk.get_id())[0],
0][0]
trk.update(
bboxes[d],
f_vecs[d]) ## for dlib re-intialize the trackers ?!
# create and initialise new trackers for unmatched detections
for i in unmatched_dets:
trk = KNNTracker(bboxes[i], f_vecs[i])
self.trackers.append(trk)
i = len(self.trackers)
ret = []
for trk in reversed(self.trackers):
d = trk.get_bbox()
if (trk.get_hit_streak() >=
self.min_hits): # or self.frame_count <= self.min_hits):
ret.append(
np.concatenate(
([int(i) for i in d], [trk.get_id()])).reshape(
1, -1)) # +1 as MOT benchmark requires positive
i -= 1
# remove dead tracklet
if (trk.get_time_since_update() > self.max_age):
self.trackers.pop(i)
if (len(ret) > 0):
inference.set_result(np.concatenate(ret))
else:
inference.set_result(np.empty((0, 5)))
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
# self.__encoder = ResNet50ExtractorAPI("", True)
self.__encoder = MarsExtractorAPI(flush_pipe_on_read=True)
self.__encoder.use_session_runner(session_runner)
self.__enc_in_pipe = self.__encoder.get_in_pipe()
self.__enc_out_pipe = self.__encoder.get_out_pipe()
self.__encoder.run()
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__job(inference)
else:
self.__in_pipe.wait()
def __job(self, inference):
self.__enc_in_pipe.push(
Inference(inference.get_data(),
meta_dict={'inference': inference},
return_pipe=self.__out_pipe))
class TFObjectDetectionAPI:
@staticmethod
def __get_dir_path():
return dirname(realpath(__file__))
@staticmethod
def __download_model(model_path, download_base, model_file):
print("downloading model...", model_path)
try:
os.mkdir(model_path)
except:
pass
opener = urllib.request.URLopener()
opener.retrieve(download_base + model_file, model_path + model_file)
print("finished downloading. extracting...")
tar_file = tarfile.open(model_path + model_file)
for file in tar_file.getmembers():
file_name = os.path.basename(file.name)
if 'frozen_inference_graph.pb' in file_name:
tar_file.extract(file, model_path)
print("finished extracting.")
@staticmethod
def __fetch_model_path(model_name):
dir_path = pretrained_path.get()
model_path = dir_path + '/'
model_file = model_name + '.tar.gz'
download_base = 'http://download.tensorflow.org/models/object_detection/'
path_to_frozen_graph = model_path + model_name + '/frozen_inference_graph.pb'
if not path.exists(path_to_frozen_graph):
TFObjectDetectionAPI.__download_model(model_path, download_base, model_file)
return path_to_frozen_graph
@staticmethod
def __fetch_category_indices():
dir_path = TFObjectDetectionAPI.__get_dir_path()
path_to_labels = os.path.join(data_path.get(), 'mscoco_label_map.pbtxt')
class_count = 90
label_map = label_map_util.load_labelmap(path_to_labels)
categories = label_map_util.convert_label_map_to_categories(label_map,
max_num_classes=class_count,
use_display_name=True)
category_index = label_map_util.create_category_index(categories)
category_dict = {}
for item in category_index.values():
category_dict[item['id']] = item['name']
category_dict[item['name']] = item['id']
return category_index, category_dict
def __init__(self, model_name=PRETRAINED_ssd_mobilenet_v1_coco_2017_11_17, image_shape=None,
graph_prefix=None, flush_pipe_on_read=False):
self.__category_index, self.__category_dict = self.__fetch_category_indices()
self.__path_to_frozen_graph = self.__fetch_model_path(model_name)
self.__flush_pipe_on_read = flush_pipe_on_read
self.__image_shape = image_shape
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
self.__run_session_on_thread = False
if not graph_prefix:
self.__graph_prefix = ''
else:
self.__graph_prefix = graph_prefix + '/'
def __in_pipe_process(self, inference):
image = inference.get_input()
data = np.expand_dims(image, axis=0)
inference.set_data(data)
return inference
def __out_pipe_process(self, result):
result, inference = result
num_detections = int(result['num_detections'][0])
detection_classes = result['detection_classes'][0][:num_detections].astype(np.uint8)
detection_boxes = result['detection_boxes'][0][:num_detections]
detection_scores = result['detection_scores'][0][:num_detections]
if 'detection_masks' in result:
detection_masks = result['detection_masks'][0][:num_detections]
else:
detection_masks = None
result = InferedDetections(inference.get_input(), num_detections, detection_boxes, detection_classes,
detection_scores,
masks=detection_masks, is_normalized=True, get_category_fnc=self.get_category,
annotator=self.annotate)
inference.set_result(result)
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_threading(self, run_on_thread=True):
self.__run_session_on_thread = run_on_thread
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
self.__tf_sess = session_runner.get_session()
with self.__tf_sess.graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.__path_to_frozen_graph, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name=self.__graph_prefix)
tf_default_graph = self.__tf_sess.graph
self.__image_tensor = tf_default_graph.get_tensor_by_name(self.__graph_prefix + 'image_tensor:0')
tensor_names = {output.name for op in tf_default_graph.get_operations() for output in op.outputs}
self.__tensor_dict = {}
for key in ['num_detections', 'detection_boxes', 'detection_classes', 'detection_scores', 'detection_masks']:
tensor_name = self.__graph_prefix + key + ':0'
if tensor_name in tensor_names:
self.__tensor_dict[key] = tf_default_graph.get_tensor_by_name(
tensor_name)
if 'detection_masks' in self.__tensor_dict:
# The following processing is only for single image
detection_boxes = tf.squeeze(self.__tensor_dict['detection_boxes'], [0])
detection_masks = tf.squeeze(self.__tensor_dict['detection_masks'], [0])
# Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
real_num_detection = tf.cast(self.__tensor_dict['num_detections'][0], tf.int32)
detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
detection_masks, detection_boxes, self.__image_shape[0], self.__image_shape[1])
detection_masks_reframed = tf.cast(
tf.greater(detection_masks_reframed, 0.5), tf.uint8)
# Follow the convention by adding back the batch dimension
self.__tensor_dict['detection_masks'] = tf.expand_dims(
detection_masks_reframed, 0)
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__session_runner.get_in_pipe().push(
SessionRunnable(self.__job, inference, run_on_thread=self.__run_session_on_thread))
else:
self.__in_pipe.wait()
def __job(self, inference):
self.__out_pipe.push(
(self.__tf_sess.run(self.__tensor_dict, feed_dict={self.__image_tensor: inference.get_data()}), inference))
def get_category(self, category):
return self.__category_dict[category]
@staticmethod
def annotate(inference):
annotated = inference.image.copy()
vis_util.visualize_boxes_and_labels_on_image_array(
annotated,
inference.get_boxes_tlbr(),
inference.get_classes().astype(np.int32),
inference.get_scores(),
TFObjectDetectionAPI.__fetch_category_indices()[0],
instance_masks=inference.get_masks(),
use_normalized_coordinates=True,
line_thickness=1)
return annotated
class YOLO(object):
def __init__(self,
input_size,
labels,
max_box_per_image,
anchors,
flush_pipe_on_read=False):
self.__flush_pipe_on_read = flush_pipe_on_read
self.__model_path = model_path.get()
self.__model = None
self.__input_size = input_size
self.__labels = list(labels)
self.__nb_class = len(self.__labels)
self.__nb_box = len(anchors) // 2
self.__class_wt = np.ones(self.__nb_class, dtype='float32')
self.__anchors = anchors
self.__max_box_per_image = max_box_per_image
self.__session_runner = None
self.__tf_sess = None
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
self.__run_session_on_thread = False
def use_threading(self, run_on_thread=True):
self.__run_session_on_thread = run_on_thread
def __normalize(self, image):
return image / 255.
def __preprocess(self, original, warp_points):
img_shape = original.shape
test1 = np.array(warp_points['point_set_1'], dtype="float32")
test2 = np.array(warp_points['point_set_2'], dtype="float32")
# print(test2)
# points = [warp_points['point_set_2'][0],warp_points['point_set_2'][2]]
# print(points)
# np.array([[0, 0], [img_shape[0], 0], [img_shape[0], img_shape[1]], [0, img_shape[1]]], dtype=np.float32)
M = cv2.getPerspectiveTransform(test1, test2)
warped_image = cv2.warpPerspective(original,
M, (img_shape[1], img_shape[0]),
flags=cv2.INTER_LINEAR)
# warped_image = warped_image[points[0][1]:points[1][1],points[0][0]:points[1][0],:]
return warped_image
def __prepare_data(self, image):
image_h, image_w, _ = image.shape
preprocessed = cv2.resize(image,
(self.__input_size, self.__input_size))
preprocessed = self.__normalize(preprocessed)
preprocessed = preprocessed[:, :, ::-1]
preprocessed = np.expand_dims(preprocessed, 0)
return preprocessed
def __in_pipe_process(self, inference):
image = inference.get_input()
warp_points = inference.get_meta_dict()['warp_points']
dummy_array = np.zeros((1, 1, 1, 1, self.__max_box_per_image, 4))
warped_image = self.__preprocess(image, warp_points)
data = (self.__prepare_data(warped_image), dummy_array)
inference.set_data(data)
inference.get_meta_dict()['warped_image'] = warped_image
return inference
def __out_pipe_process(self, result):
result, inference = result
image = inference.get_meta_dict()['warped_image']
i_dets = RetailInference(image, self.__labels)
i_dets.decode_netout(result, self.__anchors, self.__nb_class)
inference.set_result(i_dets)
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
K.set_session(self.__session_runner.get_session())
self.__tf_sess = K.get_session()
with self.__tf_sess.as_default():
with self.__tf_sess.graph.as_default():
self.__model = load_model(self.__model_path + '/check1.h5',
custom_objects={"tf": tf})
print("Successful load of yolo")
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__session_runner.get_in_pipe().push(
SessionRunnable(
self.__job,
inference,
run_on_thread=self.__run_session_on_thread))
else:
self.__in_pipe.pull_wait()
def __job(self, inference):
data = inference.get_data()
self.__out_pipe.push((self.__model.predict([data[0],
data[1]])[0], inference))
class OFISTObjectTrackingAPI:
def __init__(self, max_age=10000, min_hits=5, flush_pipe_on_read=False, use_detection_mask=False, conf_path=None):
self.max_age = max_age
self.min_hits = min_hits
self.trackers = []
self.frame_count = 0
self.__bg_frame = None
self.__bg_gray = None
self.__conf_path = conf_path
self.__flush_pipe_on_read = flush_pipe_on_read
# self.__feature_dim = (128)
# self.__image_shape = (128, 64, 3)
self.__thread = None
self.__in_pipe = Pipe(self.__in_pipe_process)
self.__out_pipe = Pipe(self.__out_pipe_process)
self.__use_detection_mask = use_detection_mask
self.__zones = None
if self.__conf_path is not None:
self.__zones = Zone.create_zones_from_conf(self.__conf_path)
number = 0
def __extract_image_patch(self, image, bbox, patch_shape):
sx, sy, ex, ey = np.array(bbox).astype(np.int)
dx = ex - sx
dy = ey - sy
mx = int((sx + ex) / 2)
my = int((sy + ey) / 2)
dx = int(min(40, dx / 2))
dy = int(min(50, dy / 2))
# image = image[sy:my + dy, mx - dx:mx + dx]
# image = ImageEnhancer.gaussian_blurr(image, sigma=1.75)
# image = ImageEnhancer.lab_enhancement(image, l=0.75)
# image = ImageEnhancer.hsv_enhancement(image, s=10, v=5)
# image = ImageEnhancer.hls_enhancement(image, l=2)
# image = ImageEnhancer.lab_enhancement(image, l=1.25)
# image = ImageEnhancer.gamma_correction(image, gamma=3)
# image = ImageEnhancer.gaussian_blurr(image, sigma=1.1)
# dx = int(min(60, dx / 2))
# dy = int(min(90, dy / 2))
image = image[sy:my + dy, mx - dx:mx + dx]
image = ImageEnhancer.gaussian_blurr(image, sigma=1.75)
image = ImageEnhancer.lab_enhancement(image, l=0.125)
image = ImageEnhancer.hsv_enhancement(image, s=5, v=5)
image = ImageEnhancer.hls_enhancement(image, l=2)
# image = ImageEnhancer.lab_enhancement(image, l=1)
# image = ImageEnhancer.gamma_correction(image, gamma=3)
image = ImageEnhancer.gaussian_blurr(image, sigma=1.25)
# image = image[sy:ey, sx:ex]
# image = ImageEnhancer.gaussian_blurr(image, sigma=2)
# image = ImageEnhancer.lab_enhancement(image, l=0.75)
# image = ImageEnhancer.hsv_enhancement(image, s=3, v=2)
# image = ImageEnhancer.lab_enhancement(image, l=1.25)
# image = ImageEnhancer.gamma_correction(image, gamma=3)
# image = ImageEnhancer.preprocess_retinex(image)
image = cv2.resize(image, tuple(patch_shape[::-1]))
return image
def __in_pipe_process(self, inference):
i_dets = inference.get_input()
frame = i_dets.get_image()
classes = i_dets.get_classes()
boxes = i_dets.get_boxes_tlbr(normalized=False)
masks = i_dets.get_masks()
bboxes = []
scores = i_dets.get_scores()
for i in range(len(classes)):
if classes[i] == i_dets.get_category('person') and scores[i] > .95:
bboxes.append([boxes[i][1], boxes[i][0], boxes[i][3], boxes[i][2]])
patches = [0 for x in bboxes]
# flips = [0 for x in bboxes]
threads = []
for i in range(len(bboxes)):
box = bboxes[i]
if self.__use_detection_mask:
mask = masks[i]
mask = np.stack((mask, mask, mask), axis=2)
image = np.multiply(frame, mask)
else:
image = frame
def exec(patches, index):
index = i
patch = self.__extract_image_patch(image, box, self.__image_shape[:2])
if patch is None:
print("WARNING: Failed to extract image patch: %s." % str(box))
patch = np.random.uniform(0., 255., self.__image_shape).astype(np.uint8)
if bool(random.getrandbits(1)):
patches[index] = patch
else:
patches[index] = cv2.flip(patch,1)
threads.append(Thread(target=exec, args=(patches,i, )))
threads[-1].start()
for thread in threads:
thread.join()
inference.set_data(patches)
inference.get_meta_dict()['bboxes'] = bboxes
return inference
def __out_pipe_process(self, inference):
f_vecs = inference.get_result()
inference = inference.get_meta_dict()['inference']
bboxes = inference.get_meta_dict()['bboxes']
patches = inference.get_data()
self.frame_count += 1
matched, unmatched_dets, unmatched_trks = Tracker.associate_detections_to_trackers(f_vecs, self.trackers,
bboxes)
if bboxes:
for trk in self.trackers:
if (trk.get_id() not in unmatched_trks):
d = matched[np.where(matched[:, 1] == trk.get_id())[0], 0][0]
trk.update(bboxes[d], f_vecs[d], patches[d])
for i in unmatched_dets:
trk = Tracker(bboxes[i], f_vecs[i], patches[i], self.frame_count, zones=self.__zones)
self.trackers.append(trk)
i = len(self.trackers)
ret = []
trails = {}
for trk in reversed(self.trackers):
if (trk.get_hit_streak() >= self.min_hits): # or self.frame_count <= self.min_hits):
ret.append(trk)
i -= 1
if (trk.get_time_since_update() > self.max_age):
self.trackers.pop(i)
if self.frame_count - trk.get_creation_time() >= 30 and not trk.is_confident():
self.trackers.pop(i)
trails[trk.get_id()] = trk.get_trail()
#
inference.get_meta_dict()['trails'] = trails
if (len(ret) > 0):
inference.set_result(ret)
else:
inference.set_result(np.empty((0, 5)))
return inference
def get_in_pipe(self):
return self.__in_pipe
def get_out_pipe(self):
return self.__out_pipe
def get_zones(self):
return self.__zones
def use_session_runner(self, session_runner):
self.__session_runner = session_runner
# self.__encoder = ResNet50ExtractorAPI("", True)
self.__encoder = MarsExtractorAPI(flush_pipe_on_read=True)
self.__encoder.use_session_runner(session_runner)
self.__image_shape = self.__encoder.get_input_shape()
self.__enc_in_pipe = self.__encoder.get_in_pipe()
self.__enc_out_pipe = self.__encoder.get_out_pipe()
self.__encoder.run()
def run(self):
if self.__thread is None:
self.__thread = Thread(target=self.__run)
self.__thread.start()
def __run(self):
while self.__thread:
if self.__in_pipe.is_closed():
self.__enc_in_pipe.close()
self.__out_pipe.close()
return
ret, inference = self.__in_pipe.pull(self.__flush_pipe_on_read)
if ret:
self.__job(inference)
else:
self.__in_pipe.wait()
def __job(self, inference):
self.__enc_in_pipe.push(
Inference(inference.get_data(), meta_dict={'inference': inference}, return_pipe=self.__out_pipe))