def run_benchmark(model):
"""Measures training time for given model with random data.
Args:
model: string, file name of the input model.
Returns:
float, training time in ms.
"""
engine = ImprintingEngine(test_utils.test_data_path(model),
keep_classes=False)
extractor = make_interpreter(engine.serialize_extractor_model(),
device=':0')
extractor.allocate_tensors()
width, height = common.input_size(extractor)
np.random.seed(12345)
# 10 Categories, each has 20 images.
data_by_category = collections.defaultdict(list)
for i in range(10):
for _ in range(20):
data_by_category[i].append(
np.random.randint(0, 256, (height, width, 3), dtype=np.uint8))
delegate = load_edgetpu_delegate({'device': ':0'})
inference_time = 0.
for class_id, tensors in enumerate(data_by_category.values()):
for tensor in tensors:
common.set_input(extractor, tensor)
extractor.invoke()
engine.train(classify.get_scores(extractor), class_id=class_id)
start = time.perf_counter()
interpreter = tflite.Interpreter(
model_content=engine.serialize_model(),
experimental_delegates=[delegate])
interpreter.allocate_tensors()
common.set_input(interpreter, tensors[0])
interpreter.invoke()
classify.get_classes(interpreter, top_k=3)
inference_time += (time.perf_counter() - start) * 1000
print('Model: %s' % model)
print('Inference time: %.2fms' % inference_time)
return inference_time
def extract_embeddings(image_paths, interpreter):
"""Uses model to process images as embeddings.
Reads image, resizes and feeds to model to get feature embeddings. Original
image is discarded to keep maximum memory consumption low.
Args:
image_paths: ndarray, represents a list of image paths.
interpreter: TFLite interpreter, wraps embedding extractor model.
Returns:
ndarray of length image_paths.shape[0] of embeddings.
"""
input_size = common.input_size(interpreter)
feature_dim = classify.num_classes(interpreter)
embeddings = np.empty((len(image_paths), feature_dim), dtype=np.float32)
for idx, path in enumerate(image_paths):
with test_image(path) as img:
common.set_input(interpreter, img.resize(input_size, Image.NEAREST))
interpreter.invoke()
embeddings[idx, :] = classify.get_scores(interpreter)
return embeddings
def run_benchmark(model):
"""Measures training time for given model with random data.
Args:
model: string, file name of the input model.
Returns:
float, training time in ms.
"""
engine = ImprintingEngine(
test_utils.test_data_path(model), keep_classes=False)
extractor = make_interpreter(engine.serialize_extractor_model())
extractor.allocate_tensors()
width, height = common.input_size(extractor)
np.random.seed(12345)
# 10 Categories, each has 20 images.
data_by_category = collections.defaultdict(list)
for i in range(10):
for _ in range(20):
data_by_category[i].append(
np.random.randint(0, 256, (height, width, 3), dtype=np.uint8))
start = time.perf_counter()
for class_id, tensors in enumerate(data_by_category.values()):
for tensor in tensors:
common.set_input(extractor, tensor)
extractor.invoke()
engine.train(classify.get_scores(extractor), class_id=class_id)
engine.serialize_model()
training_time = (time.perf_counter() - start) * 1000
print('Model: %s' % model)
print('Training time: %.2fms' % training_time)
return training_time
def detect_and_classify_faces(detector,
classifier,
image,
threshold,
padding=10):
predictions = []
boxes = []
faces = []
height, width, _ = image.shape
detector_target_size = common.input_size(detector)
classifier_target_size = common.input_size(classifier)
scale_x, scale_y = width / detector_target_size[
0], height / detector_target_size[1]
resized_image = cv2.resize(image, detector_target_size)
run_inference(detector, resized_image.tobytes())
objects = detect.get_objects(detector, threshold)
for object in objects:
bbox = object.bbox.scale(scale_x, scale_y)
startX, startY = int(bbox.xmin - padding), int(bbox.ymin - padding)
endX, endY = int(bbox.xmax + padding), int(bbox.ymax + padding)
# ensure the bounding boxes fall within the dimensions of the image
(startX, startY) = (max(1, startX), max(1, startY))
(endX, endY) = (min(width - 1, endX), min(height - 1, endY))
boxes.append((startX, startY, endX, endY))
face = image[startY:endY, startX:endX]
face = cv2.resize(face, classifier_target_size)
faces.append(face)
for face in faces:
run_inference(classifier, face.tobytes())
prediction = classify.get_scores(classifier)
predictions.append(prediction)
return (boxes, predictions)
def _train_and_test(self, model_path, train_points, test_points,
keep_classes):
# Train.
engine = ImprintingEngine(model_path, keep_classes)
extractor = make_interpreter(engine.serialize_extractor_model(),
device=':0')
extractor.allocate_tensors()
for point in train_points:
for image in point.images:
with test_utils.test_image('imprinting', image) as img:
set_input(extractor, img)
extractor.invoke()
embedding = classify.get_scores(extractor)
self.assertEqual(len(embedding), engine.embedding_dim)
engine.train(embedding, point.class_id)
# Test.
trained_model = engine.serialize_model()
classifier = make_interpreter(trained_model, device=':0')
classifier.allocate_tensors()
self.assertEqual(len(classifier.get_output_details()), 1)
if not keep_classes:
self.assertEqual(len(train_points),
classify.num_classes(classifier))
for point in test_points:
with test_utils.test_image('imprinting', point.image) as img:
set_input(classifier, img)
classifier.invoke()
top = classify.get_classes(classifier, top_k=1)[0]
self.assertEqual(top.id, point.class_id)
self.assertGreater(top.score, point.score)
return trained_model
def train(capture_dir, labels, model, out_model):
engine = ImprintingEngine(model, keep_classes=False)
extractor = make_interpreter(engine.serialize_extractor_model(),
device=':0')
extractor.allocate_tensors()
for class_id in sorted(labels):
class_name = labels[class_id]
print('\nClass: %s (id=%d)' % (class_name, class_id))
class_capture_dir = os.path.join(capture_dir, class_name)
for img in os.listdir(class_capture_dir):
imgpath = os.path.join(class_capture_dir, img)
common.set_input(extractor,
read_image(imgpath, common.input_size(extractor)))
extractor.invoke()
embedding = classify.get_scores(extractor)
print(' %s => %s' % (imgpath, embedding))
engine.train(embedding, class_id)
with open(out_model, 'wb') as f:
f.write(engine.serialize_model())
print('\nTrained model was saved to %s' % out_model)
def _transfer_learn_and_evaluate(self, model_path, keep_classes, dataset_path,
test_ratio, top_k_range):
"""Transfer-learns with given params and returns the evaluation result.
Args:
model_path: string, path of the base model.
keep_classes: bool, whether to keep base model classes.
dataset_path: string, path to the directory of dataset. The images should
be put under sub-directory named by category.
test_ratio: float, the ratio of images used for test.
top_k_range: int, top_k range to be evaluated. The function will return
accuracy from top 1 to top k.
Returns:
list of float numbers.
"""
engine = ImprintingEngine(model_path, keep_classes)
extractor = make_interpreter(engine.serialize_extractor_model())
extractor.allocate_tensors()
num_classes = engine.num_classes
print('--------------- Parsing dataset ----------------')
print('Dataset path:', dataset_path)
# train in fixed order to ensure the same evaluation result.
train_set, test_set = test_utils.prepare_data_set_from_directory(
dataset_path, test_ratio, True)
print('Image list successfully parsed! Number of Categories = ',
len(train_set))
print('--------------- Processing training data ----------------')
print('This process may take more than 30 seconds.')
train_input = []
labels_map = {}
for class_id, (category, image_list) in enumerate(train_set.items()):
print('Processing {} ({} images)'.format(category, len(image_list)))
train_input.append(
[os.path.join(dataset_path, category, image) for image in image_list])
labels_map[num_classes + class_id] = category
# train
print('---------------- Start training -----------------')
size = common.input_size(extractor)
for class_id, images in enumerate(train_input):
for image in images:
with test_image(image) as img:
common.set_input(extractor, img.resize(size, Image.NEAREST))
extractor.invoke()
engine.train(classify.get_scores(extractor),
class_id=num_classes + class_id)
print('---------------- Training finished -----------------')
with test_utils.temporary_file(suffix='.tflite') as output_model_path:
output_model_path.write(engine.serialize_model())
# Evaluate
print('---------------- Start evaluating -----------------')
classifier = make_interpreter(output_model_path.name)
classifier.allocate_tensors()
# top[i] represents number of top (i+1) correct inference.
top_k_correct_count = [0] * top_k_range
image_num = 0
for category, image_list in test_set.items():
n = len(image_list)
print('Evaluating {} ({} images)'.format(category, n))
for image_name in image_list:
with test_image(os.path.join(dataset_path, category,
image_name)) as img:
# Set threshold as a negative number to ensure we get top k
# candidates even if its score is 0.
size = common.input_size(classifier)
common.set_input(classifier, img.resize(size, Image.NEAREST))
classifier.invoke()
candidates = classify.get_classes(classifier, top_k=top_k_range)
for i in range(len(candidates)):
candidate = candidates[i]
if candidate.id in labels_map and \
labels_map[candidate.id] == category:
top_k_correct_count[i] += 1
break
image_num += n
for i in range(1, top_k_range):
top_k_correct_count[i] += top_k_correct_count[i - 1]
return [top_k_correct_count[i] / image_num for i in range(top_k_range)]
def main():
args = _parse_args()
engine = ImprintingEngine(args.model_path, keep_classes=args.keep_classes)
extractor = make_interpreter(engine.serialize_extractor_model(),
device=':0')
extractor.allocate_tensors()
shape = common.input_size(extractor)
print('--------------- Parsing data set -----------------')
print('Dataset path:', args.data)
train_set, test_set = _read_data(args.data, args.test_ratio)
print('Image list successfully parsed! Category Num = ', len(train_set))
print('---------------- Processing training data ----------------')
print('This process may take more than 30 seconds.')
train_input = []
labels_map = {}
for class_id, (category, image_list) in enumerate(train_set.items()):
print('Processing category:', category)
train_input.append(
_prepare_images(image_list, os.path.join(args.data, category),
shape))
labels_map[class_id] = category
print('---------------- Start training -----------------')
num_classes = engine.num_classes
for class_id, tensors in enumerate(train_input):
for tensor in tensors:
common.set_input(extractor, tensor)
extractor.invoke()
embedding = classify.get_scores(extractor)
engine.train(embedding, class_id=num_classes + class_id)
print('---------------- Training finished! -----------------')
with open(args.output, 'wb') as f:
f.write(engine.serialize_model())
print('Model saved as : ', args.output)
_save_labels(labels_map, args.output)
print('------------------ Start evaluating ------------------')
interpreter = make_interpreter(args.output)
interpreter.allocate_tensors()
size = common.input_size(interpreter)
top_k = 5
correct = [0] * top_k
wrong = [0] * top_k
for category, image_list in test_set.items():
print('Evaluating category [', category, ']')
for img_name in image_list:
img = Image.open(os.path.join(args.data, category,
img_name)).resize(
size, Image.NEAREST)
common.set_input(interpreter, img)
interpreter.invoke()
candidates = classify.get_classes(interpreter,
top_k,
score_threshold=0.1)
recognized = False
for i in range(top_k):
if i < len(candidates) and labels_map[
candidates[i].id] == category:
recognized = True
if recognized:
correct[i] = correct[i] + 1
else:
wrong[i] = wrong[i] + 1
print('---------------- Evaluation result -----------------')
for i in range(top_k):
print('Top {} : {:.0%}'.format(i + 1,
correct[i] / (correct[i] + wrong[i])))
