try:
    import hugsvision
except:
    !pip install -q hugsvision
    import hugsvision
    
print(hugsvision.__version__)

from hugsvision.dataio.VisionDataset import VisionDataset

train, test, id2label, label2id = VisionDataset.fromImageFolder(
	"./data/",
	test_ratio   = 0.15,
	balanced     = True,
	augmentation = True,
)

huggingface_model = 'google/vit-base-patch16-224-in21k'


from hugsvision.nnet.VisionClassifierTrainer import VisionClassifierTrainer
from transformers import ViTFeatureExtractor, ViTForImageClassification

trainer = VisionClassifierTrainer(
	model_name   = "MyKvasirV2Model",
	train      	 = train,
	test      	 = test,
	output_dir   = "./out/",
	max_epochs   = 1,
	batch_size   = 32, # On RTX 2080 Ti
    lr 		     = 2e-5,
	fp16	     = True,
	model = ViTForImageClassification.from_pretrained(
	    huggingface_model,
	    num_labels = len(label2id),
	    label2id   = label2id,
	    id2label   = id2label
	),
	feature_extractor = ViTFeatureExtractor.from_pretrained(
		huggingface_model,
	),
)

hyp, ref = trainer.evaluate_f1_score()

import os.path
from transformers import ViTFeatureExtractor, ViTForImageClassification
from hugsvision.inference.VisionClassifierInference import VisionClassifierInference

path = "./out/MyKvasirV2Model/20_2021-08-20-01-46-44/model/"
img  = "../../../samples/kvasir_v2/dyed-lifted-polyps.jpg"

classifier = VisionClassifierInference(
    feature_extractor = ViTFeatureExtractor.from_pretrained(path),
    model = ViTForImageClassification.from_pretrained(path),
)

label = classifier.predict(img_path=img)
print("Predicted class:", label)