def test_debug_info(self):
engine = BasicEngine(
test_utils.test_data_path('mobilenet_v1_1.0_224_quant.tflite'))
# Check model's input format.
input_tensor_shape = engine.get_input_tensor_shape()
self.assertListEqual([1, 224, 224, 3], input_tensor_shape.tolist())
self.assertEqual(224 * 224 * 3, engine.required_input_array_size())
# Check model's output.
output_tensors_sizes = engine.get_all_output_tensors_sizes()
self.assertListEqual([1001], output_tensors_sizes.tolist())
self.assertEqual(1, engine.get_num_of_output_tensors())
self.assertEqual(1001, engine.get_output_tensor_size(0))
self.assertEqual(1001, engine.total_output_array_size())
# Check SSD model.
ssd_engine = BasicEngine(
test_utils.test_data_path(
'mobilenet_ssd_v1_coco_quant_postprocess.tflite'))
# Check model's input format.
input_tensor_shape = ssd_engine.get_input_tensor_shape()
self.assertListEqual([1, 300, 300, 3], input_tensor_shape.tolist())
self.assertEqual(300 * 300 * 3, ssd_engine.required_input_array_size())
# Check model's output.
output_tensors_sizes = ssd_engine.get_all_output_tensors_sizes()
self.assertListEqual([80, 20, 20, 1], output_tensors_sizes.tolist())
self.assertEqual(4, ssd_engine.get_num_of_output_tensors())
self.assertEqual(80, ssd_engine.get_output_tensor_size(0))
self.assertEqual(20, ssd_engine.get_output_tensor_size(1))
self.assertEqual(20, ssd_engine.get_output_tensor_size(2))
self.assertEqual(1, ssd_engine.get_output_tensor_size(3))
self.assertEqual(121, ssd_engine.total_output_array_size())
class HeadPoseEstimator:
def __init__(self, model_head_pose):
# Init Head Pose Estimator
self.devices = edgetpu_utils.ListEdgeTpuPaths(edgetpu_utils.EDGE_TPU_STATE_UNASSIGNED)
self.engine = BasicEngine(model_path=model_head_pose, device_path=self.devices[0])
self.model_height = self.engine.get_input_tensor_shape()[1]
self.model_width = self.engine.get_input_tensor_shape()[2]
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--model',
help='Path of the segmentation model.',
required=True)
parser.add_argument(
'--input', help='File path of the input image.', required=True)
parser.add_argument('--output', help='File path of the output image.')
parser.add_argument(
'--keep_aspect_ratio',
dest='keep_aspect_ratio',
action='store_true',
help=(
'keep the image aspect ratio when down-sampling the image by adding '
'black pixel padding (zeros) on bottom or right. '
'By default the image is resized and reshaped without cropping. This '
'option should be the same as what is applied on input images during '
'model training. Otherwise the accuracy may be affected and the '
'bounding box of detection result may be stretched.'))
parser.set_defaults(keep_aspect_ratio=False)
args = parser.parse_args()
if not args.output:
output_name = 'semantic_segmentation_result.jpg'
else:
output_name = args.output
# Initialize engine.
engine = BasicEngine(args.model)
_, height, width, _ = engine.get_input_tensor_shape()
# Open image.
img = Image.open(args.input)
if args.keep_aspect_ratio:
resized_img, ratio = image_processing.resampling_with_original_ratio(
img, (width, height), Image.NEAREST)
else:
resized_img = img.resize((width, height))
ratio = (1., 1.)
input_tensor = np.asarray(resized_img).flatten()
_, raw_result = engine.run_inference(input_tensor)
result = np.reshape(raw_result, (height, width))
new_width, new_height = int(width * ratio[0]), int(height * ratio[1])
# If keep_aspect_ratio, we need to remove the padding area.
result = result[:new_height, :new_width]
vis_result = label_to_color_image(result.astype(int)).astype(np.uint8)
vis_result = Image.fromarray(vis_result)
vis_img = resized_img.crop((0, 0, new_width, new_height))
# Concat resized input image and processed segmentation results.
concated_image = Image.new('RGB', (new_width*2, new_height))
concated_image.paste(vis_img, (0, 0))
concated_image.paste(vis_result, (width, 0))
concated_image.save(output_name)
print('Please check ', output_name)
def _run_benchmark_for_model(model_name):
"""Runs benchmark for given model with a random input.
Args:
model_name: string, file name of the model.
Returns:
float, average inference time.
"""
iterations = 200 if ('edgetpu' in model_name) else 20
print('Benchmark for [', model_name, ']')
print('model path = ', test_utils.test_data_path(model_name))
engine = BasicEngine(test_utils.test_data_path(model_name))
print('Shape of input tensor : ', engine.get_input_tensor_shape())
# Prepare a random generated input.
input_size = engine.required_input_array_size()
random_input = test_utils.generate_random_input(1, input_size)
# Convert it to a numpy.array.
input_data = np.array(random_input, dtype=np.uint8)
benchmark_time = timeit.timeit(
lambda: engine.run_inference(input_data),
number=iterations)
# Time consumed for each iteration (milliseconds).
time_per_inference = (benchmark_time / iterations) * 1000
print(time_per_inference, 'ms (iterations = ', iterations, ')')
return time_per_inference
def _GetRequiredShape(model_path):
"""Gets image shape required by model.
Args:
model_path: string, path of the model.
Returns:
(width, height).
"""
tmp = BasicEngine(model_path)
input_tensor = tmp.get_input_tensor_shape()
return (input_tensor[2], input_tensor[1])
def getRequiredInputShape(self):
"""
Get the required input shape for the model.
"""
basic_engine = BasicEngine(self._model_path)
input_tensor_shape = basic_engine.get_input_tensor_shape()
if (input_tensor_shape.size != 4 or input_tensor_shape[3] != 3 or
input_tensor_shape[0] != 1):
raise RuntimeError(
'Invalid input tensor shape! Expected: [1, height, width, 3]')
return (input_tensor_shape[2], input_tensor_shape[1])
def _GetRequiredShape(model_path):
"""Gets image shape required by model.
Args:
model_path: string, path of the model.
Returns:
(width, height).
"""
tmp = BasicEngine(model_path)
input_tensor = tmp.get_input_tensor_shape()
return (input_tensor[2], input_tensor[1])
def _get_input_tensor_shape(model_path):
"""Gets input tensor shape of given model.
Args:
model_path: string, path of the model.
Returns:
List of integers.
"""
tmp = BasicEngine(model_path)
shape = tmp.get_input_tensor_shape()
return shape.copy()
def _get_shape(model):
"""Gets images shape required by model.
Args:
model: string, file name of the input model.
Returns:
(width, height)
"""
basic_engine = BasicEngine(test_utils.test_data_path('imprinting', model))
_, height, width, _ = basic_engine.get_input_tensor_shape()
return (width, height)
def test_append_fully_connected_and_softmax_layer_to_model(self):
in_model_name = (
'mobilenet_v1_1.0_224_quant_embedding_extractor_edgetpu.tflite')
in_model_path = test_utils.test_data_path(in_model_name)
in_engine = BasicEngine(in_model_path)
# Generate random input tensor.
np.random.seed(12345)
input_tensor = np.random.randint(
0, 255, size=in_engine.get_input_tensor_shape(),
dtype=np.uint8).flatten()
# Set up weights, biases and FC output tensor range.
_, embedding_vector = in_engine.run_inference(input_tensor)
embedding_vector_dim = embedding_vector.shape[0]
num_classes = 10
weights = np.random.randn(embedding_vector_dim,
num_classes).astype(np.float32)
biases = np.random.randn(num_classes).astype(np.float32)
fc_output = embedding_vector.dot(weights) + biases
fc_output_min = float(np.min(fc_output))
fc_output_max = float(np.max(fc_output))
try:
# Create temporary directory, and remove this folder when test
# finishes. Otherwise, test may fail because of generated files from
# previous run.
tmp_dir = tempfile.mkdtemp()
# Append FC and softmax layers and save model.
out_model_path = os.path.join(
tmp_dir, in_model_name[:-7] + '_bp_retrained.tflite')
AppendFullyConnectedAndSoftmaxLayerToModel(
in_model_path, out_model_path,
weights.transpose().flatten(), biases, fc_output_min,
fc_output_max)
self.assertTrue(os.path.exists(out_model_path))
# Run with saved model on same input.
out_engine = BasicEngine(out_model_path)
_, result = out_engine.run_inference(input_tensor)
# Calculate expected result.
expected = np.exp(fc_output - np.max(fc_output))
expected = expected / np.sum(expected)
np.testing.assert_almost_equal(result, expected, decimal=2)
finally:
shutil.rmtree(tmp_dir)
def run_benchmark(model):
"""Returns average inference time in ms on specified model on random input."""
print('Benchmark for [%s]' % model)
print('model path = %s' % test_utils.test_data_path(model))
engine = BasicEngine(test_utils.test_data_path(model))
print('input tensor shape = %s' % engine.get_input_tensor_shape())
iterations = 200 if 'edgetpu' in model else 20
input_size = engine.required_input_array_size()
random_input = test_utils.generate_random_input(1, input_size)
input_data = np.array(random_input, dtype=np.uint8)
result = 1000 * timeit.timeit(lambda: engine.run_inference(input_data),
number=iterations) / iterations
print('%.2f ms (iterations = %d)' % (result, iterations))
return result
class DualNetworkEdgeTpu():
"""DualNetwork implementation for Google's EdgeTPU."""
def __init__(self, save_file):
self.engine = BasicEngine(save_file)
self.board_size = go.N
self.output_policy_size = self.board_size**2 + 1
input_tensor_shape = self.engine.get_input_tensor_shape()
expected_input_shape = [1, self.board_size, self.board_size, 17]
if not np.array_equal(input_tensor_shape, expected_input_shape):
raise RuntimeError(
'Invalid input tensor shape {}. Expected: {}'.format(
input_tensor_shape, expected_input_shape))
output_tensors_sizes = self.engine.get_all_output_tensors_sizes()
expected_output_tensor_sizes = [self.output_policy_size, 1]
if not np.array_equal(output_tensors_sizes,
expected_output_tensor_sizes):
raise RuntimeError(
'Invalid output tensor sizes {}. Expected: {}'.format(
output_tensors_sizes, expected_output_tensor_sizes))
def run(self, position):
"""Runs inference on a single position."""
probs, values = self.run_many([position])
return probs[0], values[0]
def run_many(self, positions):
"""Runs inference on a list of position."""
processed = list(map(features_lib.extract_features, positions))
probabilities = []
values = []
for state in processed:
assert state.shape == (self.board_size, self.board_size,
17), str(state.shape)
result = self.engine.RunInference(state.flatten())
# If needed you can get the raw inference time from the result object.
# inference_time = result[0] # ms
policy_output = result[1][0:self.output_policy_size]
value_output = result[1][-1]
probabilities.append(policy_output)
values.append(value_output)
return probabilities, values
class image_segmentation():
MODEL_PASCAL = 'models/deeplabv3_mnv2_pascal_quant_edgetpu.tflite'
MODEL_DM05 = 'models/deeplabv3_mnv2_dm05_pascal_quant_edgetpu.tflite'
def __init__(self, model=MODEL_PASCAL):
self.engine = BasicEngine(model)
_, self.width, self.height, _ = self.engine.get_input_tensor_shape()
def segment(self, img):
self.input_image = cv2.resize(img, (self.width, self.height))
_, raw_result = self.engine.run_inference(self.input_image.flatten())
result = np.reshape(raw_result, (self.height, self.width))
self.segmented_img = self.label_to_color_image(
result.astype(int)).astype(np.uint8)
def get_original_image(self):
return self.input_image
def get_segmented_image(self):
return self.segmented_img
def label_to_color_image(self, label):
colormap = np.zeros((256, 3), dtype=int)
indices = np.arange(256, dtype=int)
for shift in reversed(range(8)):
for channel in range(3):
colormap[:, channel] |= ((indices >> channel) & 1) << shift
indices >>= 3
return colormap[label]
def write_image_to_file(self, filename):
orig = Image.fromarray(self.input_image)
segm = Image.fromarray(self.segmented_img)
concated_image = Image.new('RGB', (self.width * 2, self.height))
concated_image.paste(orig, (0, 0))
concated_image.paste(segm, (self.width, 0))
concated_image.save(filename)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model",
help="File path of Tflite model.",
required=True)
parser.add_argument("--image",
help="File path of the image to be recognized.",
required=True)
# parser.add_argument(
# '--label', help='File path of label file.', required=True)
args = parser.parse_args()
# Initialize colormap
colormap = label_util.create_pascal_label_colormap()
# Read image
org_img = Image.open(args.image)
im_width, im_height = org_img.size
engine = BasicEngine(args.model)
_, height, width, _ = engine.get_input_tensor_shape()
img = org_img.resize((width, height), Image.NEAREST)
input_tensor = np.asarray(img).flatten()
latency, result = engine.RunInference(input_tensor)
seg_map = np.array(result, dtype=np.uint8)
seg_map = np.reshape(seg_map, (width, height))
seg_image = label_util.label_to_color_image(colormap, seg_map)
seg_image = Image.fromarray(seg_image).resize((im_width, im_height),
Image.NEAREST)
out_image = np.array(org_img) * 0.5 + np.array(seg_image) * 0.5
pil_img = Image.fromarray(out_image.astype(np.uint8))
pil_img.save(os.path.join(".", "save.png"))
import time
from PIL import Image, ImageDraw, ImageFont
import numpy
from edgetpu.basic.basic_engine import BasicEngine
MODEL_NAME = "model_stride/model_stride_256x256x3_edgetpu.tflite"
### Load model and prepare TPU engine
engine = BasicEngine(MODEL_NAME)
width = engine.get_input_tensor_shape()[1]
height = engine.get_input_tensor_shape()[2]
### prepara input tensor
img = Image.new('RGB', (width, height), (128, 128, 128))
draw = ImageDraw.Draw(img)
input_tensor = numpy.asarray(img).flatten()
### Run inference
start = time.time()
num_measurement = 10000
for i in range(num_measurement):
_, raw_result = engine.RunInference(input_tensor)
# time.sleep(2)
elapsed_time = time.time() - start
print("elapsed_time: {0} ".format(1000 * elapsed_time / num_measurement) +
"[msec]")
def retrain_model(props):
"""
This function is using the Imprinting technique to retrain the model by only changing the last layer.
All classes will be abandoned while training multiple users
"""
MODEL_PATH = props['classification']['default_path']
click.echo('Parsing data for retraining...')
train_set = {}
test_set = {}
for user in props['user'].keys():
image_dir = props['user'][user]['images']
images = [
f for f in os.listdir(image_dir)
if os.path.isfile(os.path.join(image_dir, f))
]
if images:
# allocate the number of images for training an validation
net_pictures = len(images)
click.echo(
click.style('We found {} pictures for {}'.format(
net_pictures, user),
fg='green'))
while True:
k = int(
click.prompt(
'How many pictures do you want for validating the training?'
))
if k > 0.25 * net_pictures:
click.echo(
click.style(
'At most 25% ({} pictures) of the training data can be used for testing the model!'
.format(int(0.25 * net_pictures)),
fg='yellow'))
elif k < 2:
click.echo(
click.style(
'At least 3 pictues must be used for testing the model!',
fg='yellow'))
else:
break
test_set[user] = images[:k]
assert test_set, 'No images to test [{}]'.format(user)
train_set[user] = images[k:]
assert train_set, 'No images to train [{}]'.format(user)
#get shape of model to retrain
tmp = BasicEngine(MODEL_PATH)
input_tensor = tmp.get_input_tensor_shape()
shape = (input_tensor[2], input_tensor[1])
#rezising pictures and creating new labels map
train_input = []
labels_map = {}
for user_id, (user, image_list) in enumerate(train_set.items()):
ret = []
for filename in image_list:
with Image.open(
os.path.join(props['user'][user]['images'],
filename)) as img:
img = img.convert('RGB')
img = img.resize(shape, Image.NEAREST)
ret.append(np.asarray(img).flatten())
train_input.append(np.array(ret))
labels_map[user_id] = user
#Train model
click.echo('Start training')
engine = ImprintingEngine(MODEL_PATH, keep_classes=False)
engine.train_all(train_input)
click.echo(click.style('Training finished!', fg='green'))
#gethering old model files
old_model = props['classification']['path']
old_labels = props['classification']['labels']
#saving new model
props['classification']['path'] = './Models/model{}.tflite'.format(''.join(
['_' + u for u in labels_map.values()]))
engine.save_model(props['classification']['path'])
#saving labels
props['classification']['labels'] = props['classification'][
'path'].replace('classification', 'labels').replace('tflite', 'json')
with open(props['classification']['labels'], 'w') as f:
json.dump(labels_map, f, indent=4)
#Evaluating how well the retrained model performed
click.echo('Start evaluation')
engine = ClassificationEngine(props['classification']['path'])
top_k = 5
correct = [0] * top_k
wrong = [0] * top_k
for user, image_list in test_set.items():
for img_name in image_list:
img = Image.open(
os.path.join(props['user'][user]['images'], img_name))
candidates = engine.classify_with_image(img,
threshold=0.1,
top_k=top_k)
recognized = False
for i in range(top_k):
if i < len(candidates) and user == labels_map[candidates[i]
[0]]:
recognized = True
if recognized:
correct[i] = correct[i] + 1
else:
wrong[i] = wrong[i] + 1
click.echo('Evaluation Results:')
for i in range(top_k):
click.echo('Top {} : {:.0%}'.format(
i + 1, correct[i] / (correct[i] + wrong[i])))
# TODO highlight with colors how well it perforemed
if not old_model == props['classification'][
'path'] and not old_labels == props['classification'][
'labels'] and (os.path.exists(old_labels)
or os.path.exists(old_model)):
if not click.confirm('Do you want to keep old models?'):
os.remove(old_model)
os.remove(old_labels)
click.echo(click.style('Old models removed.', fg='green'))
#saving properties
save_properties(props)
parser.add_argument("--deep_model", default="bisenetv2_celebamaskhq_448x448_full_integer_quant_edgetpu.tflite", help="Path of the BiSeNetV2 model.")
parser.add_argument("--usbcamno", type=int, default=0, help="USB Camera number.")
parser.add_argument('--camera_width', type=int, default=640, help='USB Camera resolution (width). (Default=640)')
parser.add_argument('--camera_height', type=int, default=480, help='USB Camera resolution (height). (Default=480)')
parser.add_argument('--vidfps', type=int, default=30, help='FPS of Video. (Default=30)')
args = parser.parse_args()
deep_model = args.deep_model
usbcamno = args.usbcamno
vidfps = args.vidfps
camera_width = args.camera_width
camera_height = args.camera_height
devices = edgetpu_utils.ListEdgeTpuPaths(edgetpu_utils.EDGE_TPU_STATE_UNASSIGNED)
engine = BasicEngine(model_path=deep_model, device_path=devices[0])
model_height = engine.get_input_tensor_shape()[1]
model_width = engine.get_input_tensor_shape()[2]
cam = cv2.VideoCapture(usbcamno)
cam.set(cv2.CAP_PROP_FPS, vidfps)
cam.set(cv2.CAP_PROP_FRAME_WIDTH, camera_width)
cam.set(cv2.CAP_PROP_FRAME_HEIGHT, camera_height)
waittime = 1
window_name = "USB Camera"
cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)
while True:
t1 = time.perf_counter()
ret, color_image = cam.read()
class EdgeTPUSemanticSegmenter(ConnectionBasedTransport):
def __init__(self, namespace='~'):
# get image_trasport before ConnectionBasedTransport subscribes ~input
self.transport_hint = rospy.get_param(namespace + 'image_transport',
'raw')
rospy.loginfo("Using transport {}".format(self.transport_hint))
super(EdgeTPUSemanticSegmenter, self).__init__()
rospack = rospkg.RosPack()
pkg_path = rospack.get_path('coral_usb')
self.bridge = CvBridge()
self.classifier_name = rospy.get_param(namespace + 'classifier_name',
rospy.get_name())
model_file = os.path.join(
pkg_path, './models/deeplabv3_mnv2_pascal_quant_edgetpu.tflite')
model_file = rospy.get_param(namespace + 'model_file', model_file)
label_file = rospy.get_param(namespace + 'label_file', None)
if model_file is not None:
self.model_file = get_filename(model_file, False)
if label_file is not None:
label_file = get_filename(label_file, False)
self.duration = rospy.get_param(namespace + 'visualize_duration', 0.1)
self.enable_visualization = rospy.get_param(
namespace + 'enable_visualization', True)
device_id = rospy.get_param(namespace + 'device_id', None)
if device_id is None:
device_path = None
else:
device_path = ListEdgeTpuPaths(EDGE_TPU_STATE_NONE)[device_id]
assigned_device_paths = ListEdgeTpuPaths(EDGE_TPU_STATE_ASSIGNED)
if device_path in assigned_device_paths:
rospy.logwarn('device {} is already assigned: {}'.format(
device_id, device_path))
self.engine = BasicEngine(self.model_file, device_path)
self.input_shape = self.engine.get_input_tensor_shape()[1:3]
if label_file is None:
self.label_names = [
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
'train', 'tvmonitor'
]
self.label_ids = list(range(len(self.label_names)))
else:
self.label_ids, self.label_names = self._load_labels(label_file)
self.namespace = namespace
self.pub_label = self.advertise(namespace + 'output/label',
Image,
queue_size=1)
# visualize timer
if self.enable_visualization:
self.lock = threading.Lock()
self.pub_image = self.advertise(namespace + 'output/image',
Image,
queue_size=1)
self.pub_image_compressed = self.advertise(
namespace + 'output/image/compressed',
CompressedImage,
queue_size=1)
self.timer = rospy.Timer(rospy.Duration(self.duration),
self.visualize_cb)
self.img = None
self.header = None
self.label = None
def start(self):
self.engine = BasicEngine(self.model_file)
self.input_shape = self.engine.get_input_tensor_shape()[1:3]
self.subscribe()
if self.enable_visualization:
self.timer = rospy.Timer(rospy.Duration(self.duration),
self.visualize_cb)
def stop(self):
self.unsubscribe()
del self.sub_image
if self.enable_visualization:
self.timer.shutdown()
del self.timer
del self.engine
def subscribe(self):
if self.transport_hint == 'compressed':
self.sub_image = rospy.Subscriber('{}/compressed'.format(
rospy.resolve_name('~input')),
CompressedImage,
self.image_cb,
queue_size=1,
buff_size=2**26)
else:
self.sub_image = rospy.Subscriber('~input',
Image,
self.image_cb,
queue_size=1,
buff_size=2**26)
def unsubscribe(self):
self.sub_image.unregister()
@property
def visualize(self):
return self.pub_image.get_num_connections() > 0 or \
self.pub_image_compressed.get_num_connections() > 0
def config_callback(self, config, level):
self.score_thresh = config.score_thresh
self.top_k = config.top_k
return config
def _load_labels(self, path):
p = re.compile(r'\s*(\d+)(.+)')
with open(path, 'r', encoding='utf-8') as f:
lines = (p.match(line).groups() for line in f.readlines())
labels = {int(num): text.strip() for num, text in lines}
return list(labels.keys()), list(labels.values())
def _segment_step(self, img):
H, W = img.shape[:2]
input_H, input_W = self.input_shape
input_tensor = cv2.resize(img, (input_W, input_H))
input_tensor = input_tensor.flatten()
_, label = self.engine.run_inference(input_tensor)
label = label.reshape(self.input_shape)
label = cv2.resize(label, (W, H), interpolation=cv2.INTER_NEAREST)
label = label.astype(np.int32)
return label
def _segment(self, img):
return self._segment_step(img)
def image_cb(self, msg):
if not hasattr(self, 'engine'):
return
if self.transport_hint == 'compressed':
np_arr = np.fromstring(msg.data, np.uint8)
img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
img = img[:, :, ::-1]
else:
img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='rgb8')
label = self._segment(img)
label_msg = self.bridge.cv2_to_imgmsg(label, '32SC1')
label_msg.header = msg.header
self.pub_label.publish(label_msg)
if self.enable_visualization:
with self.lock:
self.img = img
self.header = msg.header
self.label = label
def visualize_cb(self, event):
if (not self.visualize or self.img is None or self.header is None
or self.label is None):
return
with self.lock:
img = self.img.copy()
header = copy.deepcopy(self.header)
label = self.label.copy()
fig = plt.figure(tight_layout={'pad': 0})
ax = plt.subplot(1, 1, 1)
ax.axis('off')
ax, legend_handles = vis_semantic_segmentation(
img.transpose((2, 0, 1)),
label,
label_names=self.label_names,
alpha=0.7,
all_label_names_in_legend=True,
ax=ax)
ax.legend(handles=legend_handles, bbox_to_anchor=(1, 1), loc=2)
fig.canvas.draw()
w, h = fig.canvas.get_width_height()
vis_img = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
vis_img.shape = (h, w, 3)
fig.clf()
plt.close()
if self.pub_image.get_num_connections() > 0:
vis_msg = self.bridge.cv2_to_imgmsg(vis_img, 'rgb8')
# BUG: https://answers.ros.org/question/316362/sensor_msgsimage-generates-float-instead-of-int-with-python3/ # NOQA
vis_msg.step = int(vis_msg.step)
vis_msg.header = header
self.pub_image.publish(vis_msg)
if self.pub_image_compressed.get_num_connections() > 0:
# publish compressed http://wiki.ros.org/rospy_tutorials/Tutorials/WritingImagePublisherSubscriber # NOQA
vis_compressed_msg = CompressedImage()
vis_compressed_msg.header = header
vis_compressed_msg.format = "jpeg"
vis_img_rgb = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
vis_compressed_msg.data = np.array(
cv2.imencode('.jpg', vis_img_rgb)[1]).tostring()
self.pub_image_compressed.publish(vis_compressed_msg)
# Initialize frame rate calculation
frame_rate_calc = 1
freq = cv2.getTickFrequency()
xoff, yoff = 0, 0
# Initialize video stream
# videostream = VideoStream(resolution=(imW,imH),framerate=30).start()
# time.sleep(1)
model = "models/posenet_mobilenet_v1_075_{}_{}_quant_decoder_edgetpu.tflite".format(
int(resH) + 1,
int(resW) + 1)
KEYPOINTS = labels
engine = BasicEngine(model)
input_shape = engine.get_input_tensor_shape()
inference_size = (input_shape[2], input_shape[1])
if (input_shape.size != 4 or input_shape[3] != 3 or input_shape[0] != 1):
raise ValueError(
('Image model should have input shape [1, height, width, 3]!'
' This model has {}.'.format(self.input_shape)))
_, image_height, image_width, image_depth = engine.get_input_tensor_shape()
# The API returns all the output tensors flattened and concatenated. We
# have to figure out the boundaries from the tensor shapes & sizes.
offset = 0
#print("offset",offset)
_output_offsets = [0]
for size in engine.get_all_output_tensors_sizes():
offset += size
_output_offsets.append(offset)
class FaceEmbeddingEngine:
''' class FaceEmbeddingEngine
Purpose: generate embeddings for images of faces
'''
def __init__(self, embedding_model):
''' function constructor
Constructor for FaceEmbeddingEngine
Args:
embedding_model (FaceEmbeddingModelEnum): The model to use for generating
embeddings for face images
Returns:
None
'''
# We only want to import these modules at run-time since
# they will only be installed on certain platforms.
# pylint: disable=import-outside-toplevel, import-error
self.embedding_model = embedding_model
self.required_image_shape = get_image_dimensions_for_embedding_model(
embedding_model) + (3, ) # need 3 arrays for RGB
if self.embedding_model == FaceEmbeddingModelEnum.CELEBRITY_KERAS:
print("Using Celebrity trained Keras model for face embeddings")
from keras.models import load_model
self.face_embedding_engine = load_model(
FACE_EMBEDDING_CELEBRITY_KERAS_MODEL_PATH, compile=False)
elif self.embedding_model == FaceEmbeddingModelEnum.CELEBRITY_TFLITE:
print("Using Celebrity trained tflite model for face embeddings")
from edgetpu.basic.basic_engine import BasicEngine
self.face_embedding_engine = BasicEngine(
FACE_EMBEDDING_CELEBRITY_TFLITE_MODEL_PATH)
print("Embedding model input tensor shape: {}".format(
self.face_embedding_engine.get_input_tensor_shape()))
print("Embedding model input size: {}".format(
self.face_embedding_engine.required_input_array_size()))
else:
raise Exception(
"Invalid embedding mode method: {}".format(embedding_model))
def get_embedding_model(self):
''' function get_embedding_model
Get the embedding model being used by this instance of the FaceEmbeddingEngine
Args:
None
Returns:
The FaceEmbeddingModelEnum being used by this instance of FaceEmbeddingEngine
'''
return self.embedding_model
# get the face embedding for one face
def get_embedding(self, face_pixels):
''' function get_embedding
Generate an embedding for the given face
Args:
face_pixels (cv2 image): The image of the face to generate the
embedding for. The dimensions of the image must
match the dimensions required by the selected
embedding model.
Returns:
A numpy array with the embedding that was generated
'''
# Confirm we're using a proper sized image to generate the embedding with
if face_pixels.shape != self.required_image_shape:
raise Exception(
"Invalid shape: {} for embedding mode method: {}".format(
face_pixels.shape, self.embedding_model))
# scale pixel values
face_pixels = face_pixels.astype('float32')
# standardize pixel values across channels (global)
mean, std = face_pixels.mean(), face_pixels.std()
face_pixels = (face_pixels - mean) / std
# transform face into one sample
sample = expand_dims(face_pixels, axis=0)
# get embedding
if self.embedding_model == FaceEmbeddingModelEnum.CELEBRITY_KERAS:
embeddings = self.face_embedding_engine.predict(sample)
result = embeddings[0]
else:
sample = sample.flatten()
# normalize values to between 0 and 255 (UINT)
sample *= 255.0 / sample.max()
# convert to UNIT8
sample = sample.astype(np_uint8)
embeddings = self.face_embedding_engine.run_inference(sample)
result = embeddings[1]
return result
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model', help='File path of Tflite model.', required=True)
parser.add_argument('--width', help='Resolution width.', default=640, type=int)
parser.add_argument('--height', help='Resolution height.', default=480, type=int)
parser.add_argument('--nano', help='Works with JETSON Nao and Pi Camera.', action='store_true')
# parser.add_argument(
# '--label', help='File path of label file.', required=True)
args = parser.parse_args()
# Initialize window.
cv2.namedWindow(
WINDOW_NAME, cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE | cv2.WINDOW_KEEPRATIO
)
cv2.moveWindow(WINDOW_NAME, 100, 200)
# Initialize colormap
colormap = label_util.create_pascal_label_colormap()
# Initialize engine.
engine = BasicEngine(args.model)
_, width, height, _ = engine.get_input_tensor_shape()
if args.nano == True:
GST_STR = 'nvarguscamerasrc \
! video/x-raw(memory:NVMM), width={0:d}, height={1:d}, format=(string)NV12, framerate=(fraction)30/1 \
! nvvidconv flip-method=2 ! video/x-raw, width=(int){2:d}, height=(int){3:d}, format=(string)BGRx \
! videoconvert \
! appsink'.format(args.width, args.height, args.width, args.height)
cap = cv2.VideoCapture(GST_STR, cv2.CAP_GSTREAMER)
else:
cap = cv2.VideoCapture(0)
cap.set(3, args.width)
cap.set(4, args.height)
while(cap.isOpened()):
_, frame = cap.read()
start_ms = time.time()
# Create inpute tensor
# camera resolution => input tensor size (513, 513)
input_buf = cv2.resize(frame, (width, height))
input_buf = cv2.cvtColor(input_buf, cv2.COLOR_BGR2RGB)
input_tensor = input_buf.flatten()
# Run inference
latency, result = engine.RunInference(input_tensor)
# Create segmentation map
seg_map = np.array(result, dtype=np.uint8)
seg_map = np.reshape(seg_map, (width, height))
seg_image = label_util.label_to_color_image(colormap, seg_map)
# segmentation map resize 513, 513 => camera resolution
seg_image = cv2.resize(seg_image, (args.width, args.height))
im = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) // 2 + seg_image // 2
im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
elapsed_ms = time.time() - start_ms
# Calc fps.
fps = 1 / elapsed_ms
fps_text = '{0:.2f}ms, {1:.2f}fps'.format((elapsed_ms * 1000.0), fps)
visual.draw_caption(im, (10, 30), fps_text)
latency_text = 'RunInference latency: {0:.2f}ms'.format(latency)
visual.draw_caption(im, (10, 60), latency_text)
# Display image
cv2.imshow(WINDOW_NAME, im)
key = cv2.waitKey(10) & 0xFF
if key == ord('q'):
break
if args.nano != True:
for i in range(10):
ret, frame = cap.read()
# When everything done, release the window
cap.release()
cv2.destroyAllWindows()
class EdgeTPUSemanticSegmenter(ConnectionBasedTransport):
def __init__(self):
super(EdgeTPUSemanticSegmenter, self).__init__()
rospack = rospkg.RosPack()
pkg_path = rospack.get_path('coral_usb')
self.bridge = CvBridge()
self.classifier_name = rospy.get_param('~classifier_name',
rospy.get_name())
model_file = os.path.join(
pkg_path, './models/deeplabv3_mnv2_pascal_quant_edgetpu.tflite')
model_file = rospy.get_param('~model_file', model_file)
label_file = rospy.get_param('~label_file', None)
self.engine = BasicEngine(model_file)
self.input_shape = self.engine.get_input_tensor_shape()[1:3]
if label_file is None:
self.label_names = [
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
'train', 'tvmonitor'
]
self.label_ids = list(range(len(self.label_names)))
else:
self.label_ids, self.label_names = self._load_labels(label_file)
self.pub_label = self.advertise('~output/label', Image, queue_size=1)
self.pub_image = self.advertise('~output/image', Image, queue_size=1)
def subscribe(self):
self.sub_image = rospy.Subscriber('~input',
Image,
self.image_cb,
queue_size=1,
buff_size=2**26)
def unsubscribe(self):
self.sub_image.unregister()
@property
def visualize(self):
return self.pub_image.get_num_connections() > 0
def config_callback(self, config, level):
self.score_thresh = config.score_thresh
self.top_k = config.top_k
return config
def _load_labels(self, path):
p = re.compile(r'\s*(\d+)(.+)')
with open(path, 'r', encoding='utf-8') as f:
lines = (p.match(line).groups() for line in f.readlines())
labels = {int(num): text.strip() for num, text in lines}
return list(labels.keys()), list(labels.values())
def image_cb(self, msg):
img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='rgb8')
H, W = img.shape[:2]
input_H, input_W = self.input_shape
input_tensor = cv2.resize(img, (input_W, input_H))
input_tensor = input_tensor.flatten()
_, label = self.engine.run_inference(input_tensor)
label = label.reshape(self.input_shape)
label = cv2.resize(label, (W, H), interpolation=cv2.INTER_NEAREST)
label = label.astype(np.int32)
label_msg = self.bridge.cv2_to_imgmsg(label, '32SC1')
label_msg.header = msg.header
self.pub_label.publish(label_msg)
if self.visualize:
fig = plt.figure(tight_layout={'pad': 0})
ax = plt.subplot(1, 1, 1)
ax.axis('off')
ax, legend_handles = vis_semantic_segmentation(
img.transpose((2, 0, 1)),
label,
label_names=self.label_names,
alpha=0.7,
all_label_names_in_legend=True,
ax=ax)
ax.legend(handles=legend_handles, bbox_to_anchor=(1, 1), loc=2)
fig.canvas.draw()
w, h = fig.canvas.get_width_height()
vis_img = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
vis_img.shape = (h, w, 3)
fig.clf()
plt.close()
vis_msg = self.bridge.cv2_to_imgmsg(vis_img, 'rgb8')
# BUG: https://answers.ros.org/question/316362/sensor_msgsimage-generates-float-instead-of-int-with-python3/ # NOQA
vis_msg.step = int(vis_msg.step)
vis_msg.header = msg.header
self.pub_image.publish(vis_msg)
import numpy as np
import pandas as pd
# from edgetpu.classification.engine import ClassificationEngine
from edgetpu.basic.basic_engine import BasicEngine
from sklearn.datasets import load_iris
data = load_iris()
x_train = (data.data * 10).astype(np.float32)
y_train = pd.get_dummies(data.target).values.astype(np.float32)
eng = BasicEngine("converted_model_edgetpu.tflite")
print(eng)
print(eng.get_all_output_tensors_sizes(), eng.get_input_tensor_shape())
for i in range(150):
# inp = np.array([1,2,3]).astype(np.uint8)
inp = x_train[i].astype(np.uint8)
# inp = [30,30,30,30]
print(inp, x_train[i], y_train[i])
print(eng.RunInference(inp)[1])
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", help="File path of Tflite model.", required=True)
parser.add_argument("--width", help="Resolution width.", default=640)
parser.add_argument("--height", help="Resolution height.", default=480)
# parser.add_argument(
# '--label', help='File path of label file.', required=True)
args = parser.parse_args()
# Initialize window.
cv2.namedWindow(
WINDOW_NAME, cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE | cv2.WINDOW_KEEPRATIO
)
cv2.moveWindow(WINDOW_NAME, 100, 200)
# Initialize colormap
colormap = label_util.create_pascal_label_colormap()
# Initialize engine.
engine = BasicEngine(args.model)
resolution_width = args.width
rezolution_height = args.height
with picamera.PiCamera() as camera:
camera.resolution = (resolution_width, rezolution_height)
camera.framerate = 30
_, width, height, _ = engine.get_input_tensor_shape()
rawCapture = PiRGBArray(camera)
# allow the camera to warmup
time.sleep(0.1)
try:
for frame in camera.capture_continuous(
rawCapture, format="rgb", use_video_port=True
):
start_ms = time.time()
rawCapture.truncate(0)
image = frame.array
# Create inpute tensor
# camera resolution (640, 480) => input tensor size (513, 513)
input_buf = Image.fromarray(image)
input_buf = input_buf.resize((width, height), Image.NEAREST)
input_tensor = np.asarray(input_buf).flatten()
# Run inference
latency, result = engine.RunInference(input_tensor)
# Create segmentation map
seg_map = np.array(result, dtype=np.uint8)
seg_map = np.reshape(seg_map, (width, height))
seg_image = label_util.label_to_color_image(colormap, seg_map)
# segmentation map resize 513, 513 => camera resolution(640, 480)
seg_image = cv2.resize(seg_image, (resolution_width, rezolution_height))
out_image = image // 2 + seg_image // 2
im = cv2.cvtColor(out_image, cv2.COLOR_RGB2BGR) # display image
elapsed_ms = time.time() - start_ms
# Calc fps.
fps = 1 / elapsed_ms
fps_text = "{0:.2f}ms, {1:.2f}fps".format((elapsed_ms * 1000.0), fps)
visual.draw_caption(im, (10, 30), fps_text)
latency_text = "Runinference latency: {0:.2f}ms".format(latency)
visual.draw_caption(im, (10, 60), latency_text)
# Display image
cv2.imshow(WINDOW_NAME, im)
key = cv2.waitKey(10) & 0xFF
if key == ord("q"):
break
finally:
camera.stop_preview()
# When everything done, release the window
cv2.destroyAllWindows()
def input_tensor_size(model):
"""Returns model input tensor size."""
engine = BasicEngine(test_utils.test_data_path(model))
batch, height, width, depth = engine.get_input_tensor_shape()
return batch * height * width * depth
import matplotlib as plt
import numpy as np
import imutils
from PIL import Image
import time
from edgetpu.basic.edgetpu_utils import ListEdgeTpuPaths, EDGE_TPU_STATE_UNASSIGNED
from edgetpu.basic.basic_engine import BasicEngine
from edgetpu.utils import dataset_utils, image_processing
import collections
import math
import enum
print(ListEdgeTpuPaths(EDGE_TPU_STATE_UNASSIGNED))
engine = BasicEngine(
'posenet_mobilenet_v1_075_481_641_quant_decoder_edgetpu.tflite')
_, height, width, _ = engine.get_input_tensor_shape()
cam = cv2.VideoCapture(2)
while True:
ret, frame = cam.read()
frame = cv2.resize(frame, (width, height), cv2.INTER_NEAREST)
canvas = frame.copy()
start = time.time()
_, result = engine.run_inference(frame[:, :, ::-1].flatten())
end = time.time()
for i in range(0, len(result) - 1, 2):
cv2.circle(canvas, (int(result[i + 1]), int(result[i])), 10,
(0, 255, 0), -1)
import numpy as np
import cv2
import matplotlib.pyplot as mp
################################################################################
# %% DEFAULT SETTINGS
################################################################################
np.set_printoptions(precision=3)
################################################################################
# %% IMPORT MODEL AND MOVE TO TPU
################################################################################
engine = BasicEngine('mnist_edgetpu.tflite')
(_, xdim, ydim, zdim) = engine.get_input_tensor_shape()
################################################################################
# %% INIT SCREEN OUTPUT
################################################################################
cap = cv2.VideoCapture(0)
mp.figure()
################################################################################
# %% RUN MODEL OFF OF CAMERA ON TPU
################################################################################
while cap.isOpened():
##### GRAB IMAGE FROM CAM
class InferenceEngine:
"""Base inferencing class that wraps Tensorflow Lite"""
def __init__(self, model_path, output_shapes):
"""
Initializes InferenceEngine. Creates interpreter, allocates tensors, and grabs input and output details.
:param model_path: Path to Tensorflow Lite Model
:param output_shapes: List of tuples, each tuple containing the shape of an output tensor
"""
# load trained model
if use_tpu:
model_path = os.path.splitext(model_path) + '_edgetpu.tflite'
self.TPU_engine = BasicEngine(model_path)
else:
self.interpreter = tf.lite.Interpreter(model_path=model_path)
self.interpreter.allocate_tensors()
# Get input and output tensors
self.input_details = self.interpreter.get_input_details()
self.output_details = self.interpreter.get_output_details()
self.output_shapes = output_shapes
def get_input_shape(self):
"""
Gets the shape of the input tensor of the model
:return: Input shape as a tuple
"""
if use_tpu:
return self.TPU_engine.get_input_tensor_shape()
else:
return self.input_details[0]['shape']
def invoke(self, input_tensor, lambda_preprocess=None):
"""
Invokes the interpreter. Basically has the interpreter run the needed calculation to have the output
tensors ready.
:param input_tensor: Input Tensor (nD numpy array)
:param lambda_preprocess: Lambda function to apply to input tensors after resizing
:return: A list of output tensors (list of numpy arrays)
"""
input_shape = self.input_details[0]['shape']
resized_tensor = cv2.resize(input_tensor,
(input_shape[1], input_shape[2]))
if lambda_preprocess:
resized_tensor = lambda_preprocess(resized_tensor)
output = []
if use_tpu:
resized_tensor = resized_tensor.flatten()
_, raw_results = self.TPU_engine.RunInference(resized_tensor)
so_far = 0
for index, shape in enumerate(self.output_shapes):
output_size = self.TPU_engine.get_output_tensor_size(index)
output.append(raw_results[so_far:so_far +
output_size].reshape(shape))
so_far += output_size
else:
self.interpreter.set_tensor(self.input_details[0]['index'],
resized_tensor.reshape(input_shape))
self.interpreter.invoke()
for index, detail in enumerate(self.output_details):
output.append(
self.interpreter.get_tensor(detail['index']).reshape(
self.output_shapes[index]))
return output