def test_resize():
config = _good_config()
img_detect = TFImageDetection(**config)
_dir = os.path.dirname(os.path.abspath(__file__))
img_path = os.path.join(_dir, 'background.jpg')
image = Image.open(img_path)
orig_width = image.size[0]
assert orig_width == 1280
orig_height = image.size[1]
assert orig_height == 720
new_size = (300, 300)
new_image = img_detect.resize(image=image, desired_size=new_size)
new_width = new_image.size[0]
assert new_width == new_size[0]
new_height = new_image.size[1]
assert new_height == new_size[1]
def test_inference_init_good_config():
config = _good_config()
img_detect = TFImageDetection(**config)
assert img_detect
assert img_detect._tfengine
assert img_detect._tfengine._model_tflite_path.endswith('.tflite')
assert img_detect._tfengine._model_edgetpu_path.endswith('.tflite')
assert img_detect._tfengine.confidence_threshold == 0.654
assert img_detect._tfengine.top_k == 123
assert img_detect._tfengine.is_quantized
assert img_detect._tfengine._model_labels_path.endswith('.txt')
def test_inference_init_bad_config():
config = {
'model': {
'tflite': 'some_bad_tflite_model',
},
'labels': 'no_labels',
'top_k': 123,
'confidence_threshold': 654,
}
with pytest.raises(AssertionError):
TFImageDetection(**config)
def test_model_inputs():
"""Verify against known model inputs."""
config = _good_config()
img_detect = TFImageDetection(**config)
tfe = img_detect._tfengine
samples = tfe.input_details[0]['shape'][0]
assert samples == 1
height = tfe.input_details[0]['shape'][1]
assert height == 300
width = tfe.input_details[0]['shape'][2]
assert width == 300
colors = tfe.input_details[0]['shape'][3]
assert colors == 3
def test_model_outputs():
"""Verify against known model outputs."""
config = _good_config()
img_detect = TFImageDetection(**config)
tfe = img_detect._tfengine
assert tfe.output_details[0]['shape'][0] == 1
scores = tfe.output_details[0]['shape'][1]
assert scores == 20
assert tfe.output_details[1]['shape'][0] == 1
boxes = tfe.output_details[1]['shape'][1]
assert boxes == 20
assert tfe.output_details[2]['shape'][0] == 1
labels = tfe.output_details[2]['shape'][1]
assert labels == 20
num = tfe.output_details[3]['shape'][0]
assert num == 1
def test_inference_init_no_config():
with pytest.raises(AssertionError):
TFImageDetection()
def test_load_labels():
config = _good_config()
img_detect = TFImageDetection(**config)
labels = img_detect._labels
assert labels[0] == 'person'
assert labels[15] == 'bird'
def test_receive_next_sample():
config = _good_config()
img_detect = TFImageDetection(**config)
# no action expected from the abstract method
img_detect.receive_next_sample(image=None)
def DetectPosesInImage(self, img):
"""
Detects poses in a given image.
:Parameters:
----------
img : PIL.Image
Input Image for AI model detection.
:Returns:
-------
poses:
A list of Pose objects with keypoints and confidence scores
PIL.Image
Resized image fitting the AI model input tensor.
"""
_tensor_input_size = (self._tensor_image_width,
self._tensor_image_height)
# thumbnail is a proportionately resized image
thumbnail = TFImageDetection.thumbnail(image=img,
desired_size=_tensor_input_size)
# convert thumbnail into an image with the exact size
# as the input tensor preserving proportions by padding with a solid color as needed
template_image = TFImageDetection.resize(
image=thumbnail, desired_size=_tensor_input_size)
template_input = np.expand_dims(template_image.copy(), axis=0)
floating_model = self._tfengine.input_details[0]['dtype'] == np.float32
if floating_model:
template_input = (np.float32(template_input) - 127.5) / 127.5
self.tf_interpreter().set_tensor(
self._tfengine.input_details[0]['index'], template_input)
self.tf_interpreter().invoke()
template_output_data = self.tf_interpreter().get_tensor(
self._tfengine.output_details[0]['index'])
template_offset_data = self.tf_interpreter().get_tensor(
self._tfengine.output_details[1]['index'])
template_heatmaps = np.squeeze(template_output_data)
template_offsets = np.squeeze(template_offset_data)
kps = self.parse_output(template_heatmaps, template_offsets, 0.3)
poses = []
keypoint_dict = {}
cnt = 0
for point_i in range(kps.shape[0]):
x, y = kps[point_i, 1], kps[point_i, 0]
prob = self.sigmoid(kps[point_i, 3])
if prob > 0.60:
cnt += 1
keypoint = Keypoint(KEYPOINTS[point_i], [x, y], prob)
keypoint_dict[KEYPOINTS[point_i]] = keypoint
# draw on image and save it for debugging
draw = ImageDraw.Draw(template_image)
draw.line(((0, 0), (x, y)), fill='red')
pose_scores = cnt / 17
poses.append(Pose(keypoint_dict, pose_scores))
# DEBUG: save template_image for debugging
# DEBUG: timestr = int(time.monotonic()*1000)
# DEBUG: template_image.save(f'tmp-template-image-time-{timestr}-keypoints-{cnt}.jpg', format='JPEG')
return poses, thumbnail