in modules/edge-inference/src/inference.py [0:0]
def main():
print("==============")
print("AiTPS Software")
print("==============")
parser = argparse.ArgumentParser(description='Fuel Pump')
ipc_client = awsiot.greengrasscoreipc.connect()
# # reading the input arguemenets
parser.add_argument('--model',
help='model', default='./model')
parser.add_argument('--class_names',
help='class_names', default='./labels.txt')
parser.add_argument('--input',
help='input folder path', default='./data')
args = parser.parse_args()
class_names = load_class_names(args.class_names)
while True:
try:
start_time = datetime.utcnow()
img_path = random.choice(os.listdir(args.input))
print("Randomly selected image {}".format(img_path))
input_image = Image.open(os.path.join(args.input, img_path)).resize((224,224))
input_image.save(TEMP_IMG_PATH)
input_matrix = (np.array(input_image) / 255 - [0.485, 0.456, 0.406])/[0.229, 0.224, 0.225]
transposed = np.transpose(input_matrix, (2, 0, 1))
model = dlr.DLRModel(args.model, 'cpu', 0)
res = softmax(model.run(transposed)[0][0])
pred = np.argmax(res)
prob = res[pred]
image_bytes = None
with open(TEMP_IMG_PATH, "rb") as imageFile:
image_bytes = base64.b64encode(imageFile.read()).decode('utf8')
end_time = datetime.utcnow()
delta = end_time - start_time
total_time = delta.total_seconds() * 1000
start_time_str = start_time.strftime("%Y-%m-%d %H:%M:%S.%f")
end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S.%f")
outgoing_msg = OUTGOING_MSG_FORMAT
outgoing_msg['InferenceTotalTime'] = total_time
outgoing_msg['InferenceStartTime'] = start_time_str
outgoing_msg['InferenceEndTime'] = end_time_str
outgoing_msg['Prediction'] = class_names[pred]
outgoing_msg['Probability'] = str(prob)
outgoing_msg['Picture'] = image_bytes
publishResults(ipc_client, PUBLISH_TOPIC, outgoing_msg)
time.sleep(5)
except Exception as e:
print(e)