def main()

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)