from huggingface_hub import push_to_hub_fastai
from huggingface_hub import login
from fastai.vision.all import *

# repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
repo_id = "tobthecreator/pasta-model"


pasta_shapes = [
    "spaghetti",
    "penne",
    "fettuccine",
    "rigatoni",
    "bowtie",
    "linguine",
    "ravioli",
    "lasagna",
    "rotini",
]


if __name__ == '__main__':
    login()
    
    categories = ['bowtie', 'fettuccine', 'lasagna', 'linguine', 'not pasta', 'penne', 'ravioli', 'rigatoni', 'rotini', 'spaghetti']

    def get_y(file_path):
        parent_folder = file_path.parent.name

        if parent_folder in pasta_shapes:
            return parent_folder
        else:
            return "not pasta"
        
    learn = load_learner('model.pkl')
    push_to_hub_fastai(learner=learn, repo_id=repo_id)


    # def classify_image(img):
    #     _, _, probs = learn.predict(img)

    #     return dict(zip(categories, map(float, probs)))