Model.py

from nail_classification.inference import Inference


class Model:
    def __init__(self):
        file_paths = [hf_hub_download("lfolle/DeepNAPSIModel", f"version_{v}.ckpt",
                                use_auth_token=os.environ['DeepNAPSIModel']) for v in [10, 11, 12, 13, 14]]
        self.inference = Inference(file_paths)


    def predict(self, x):
        y_hat, uncertainty = self.inference.predict(x)
        return y_hat, uncertainty