openmv/scripts/libraries/ml/ml/model.py
Ibrahim Abdelkader 357b6f9b07
modules/py_ml: Fix model kwargs. (#2280)
* modules/py_ml: Fix model kwargs.
2024-07-14 23:44:37 +02:00

19 lines
627 B
Python

# This file is part of the OpenMV project.
#
# Copyright (c) 2024 Ibrahim Abdelkader <iabdalkader@openmv.io>
# Copyright (c) 2024 Kwabena W. Agyeman <kwagyeman@openmv.io>
#
# This work is licensed under the MIT license, see the file LICENSE for details.
import uml
import image
from ml.preprocessing import Normalization
class Model(uml.Model):
def __init__(self, *args, **kwargs):
super().__init__(*args, kwargs.get("load_to_fb", False))
def predict(self, args, **kwargs):
args = [Normalization()(x) if isinstance(x, image.Image) else x for x in args]
return super().predict(args, **kwargs)