Added tflite model & tflite python script
This commit is contained in:
parent
69e479e8a5
commit
41e2802aa0
BIN
models/converted_model.tflite
Normal file
BIN
models/converted_model.tflite
Normal file
Binary file not shown.
@ -0,0 +1,66 @@
|
||||
#!/usr/bin/python
|
||||
|
||||
"""
|
||||
Ressources :
|
||||
* https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/raspberry_pi
|
||||
"""
|
||||
|
||||
import tflite_runtime.interpreter as tflite
|
||||
import argparse
|
||||
import glob
|
||||
import time
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def set_input_tensor(interpreter, image):
|
||||
tensor_index = interpreter.get_input_details()[0]['index']
|
||||
input_tensor = interpreter.tensor(tensor_index)()[0]
|
||||
input_tensor[:, :] = image
|
||||
|
||||
def classify_image(interpreter, image):
|
||||
"""Returns a sorted array of classification results."""
|
||||
set_input_tensor(interpreter, image)
|
||||
interpreter.invoke()
|
||||
output_details = interpreter.get_output_details()[0]
|
||||
output = np.squeeze(interpreter.get_tensor(output_details['index']))
|
||||
|
||||
# If the model is quantized (uint8 data), then dequantize the results
|
||||
if output_details['dtype'] == np.uint8:
|
||||
scale, zero_point = output_details['quantization']
|
||||
output = scale * (output - zero_point)
|
||||
|
||||
ordered = np.argpartition(-output, top_k)
|
||||
return [(i, output[i]) for i in ordered[:top_k]]
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument(
|
||||
'--model', help='File path of .tflite file.', required=True)
|
||||
# parser.add_argument(
|
||||
# '--labels', help='File path of labels file.', required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
interpreter = tflite.Interpreter(args.model)
|
||||
interpreter.allocate_tensors()
|
||||
|
||||
_, height, width, _ = interpreter.get_input_details()[0]['shape']
|
||||
|
||||
CAPTURES_DIR = "/home/pi/captures"
|
||||
|
||||
image_paths = glob.glob(CAPTURES_DIR + "*.jpg")
|
||||
|
||||
for image_path in image_paths:
|
||||
im = Image.open(image_path)
|
||||
im = im.resize((width, height))
|
||||
start_time = time.time()
|
||||
results = classify_image(interpreter, im)
|
||||
elapsed_time = (time.time() - start_time) * 1000
|
||||
label_id, prob = results[0]
|
||||
|
||||
print(image_path)
|
||||
print('-' * 20)
|
||||
print(f"Prediction : {label_id} ({prob}) -- computed in {elapsed_time}ms.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
Loading…
Reference in New Issue
Block a user