From 3f55d956c52902003a6fb2b3431e65a31b772c2a Mon Sep 17 00:00:00 2001 From: iabdalkader Date: Sat, 28 Dec 2024 11:09:21 +0100 Subject: [PATCH] scripts/libraries: Load model's labels from file. Signed-off-by: iabdalkader --- scripts/libraries/ml/ml/model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/scripts/libraries/ml/ml/model.py b/scripts/libraries/ml/ml/model.py index 5ec3c9371..b42005b78 100644 --- a/scripts/libraries/ml/ml/model.py +++ b/scripts/libraries/ml/ml/model.py @@ -34,6 +34,11 @@ from ml.preprocessing import Normalization class Model(uml.Model): def __init__(self, *args, **kwargs): super().__init__(*args, kwargs.get("load_to_fb", False)) + try: + path = args[0].split(".")[0] + ".txt" + self.labels = [line.rstrip('\n') for line in open(path, "r")] + except Exception: + self.labels = None def predict(self, args, **kwargs): args = [Normalization()(x) if isinstance(x, image.Image) else x for x in args]