from datasets import load_dataset, DatasetDict, Dataset
from torchvision.transforms import Compose, ColorJitter, ToTensor
import matplotlib.pyplot as plt
import numpy as np

def main():
    # Load
    dataset = load_dataset("mnist")

    # Inspect
    # print(dataset)
    # dataset["train"][0]['image'].show()
    # dataset["test"][0]['image'].show()

    # Prepare labels mapping to be able to go from id to label
    labels = dataset["train"].features["label"].names
    label2id, id2label = {}, {}
    for i, label in enumerate(labels):
        label2id[label] = str(i)
        id2label[i] = label
    # Now check e.g. label for image 2:
    print(f"Image 2 is labelled {id2label[2]}")


def use_image_multi_class_classification():
    from transformers import AutoFeatureExtractor, AutoModelForImageClassification

    extractor = AutoFeatureExtractor.from_pretrained("autoevaluate/image-multi-class-classification")
    model = AutoModelForImageClassification.from_pretrained("autoevaluate/image-multi-class-classification")

if __name__ == '__main__':
    main()