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

import os
import os.path
from shutil import copyfile

from tqdm import tqdm

import pandas as pd

df = pd.read_csv("./train_data.csv")

img_in = "./small_train_data_set/small_train_data_set/"
img_out = "./data/"

for index, row in tqdm(df.iterrows())   :

    label = "pneumothorax" if row['target'] == 1 else "normal"

    path_in = img_in + row['file_name']
    path_out = img_out + label + "/" + row['file_name']

    # Check if the input image exist
    if not os.path.isfile(path_in):
        continue

    # Check if the output dir of the label exist
    if not os.path.isdir(img_out + label):
        os.mkdir(img_out + label)
        print("Directory for the label " + label + " created!")

    # Copy the image
    copyfile(path_in, path_out)

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 transformers import ViTFeatureExtractor, ViTForImageClassification
from hugsvision.nnet.VisionClassifierTrainer import VisionClassifierTrainer

trainer = VisionClassifierTrainer(
	model_name   = "MyPneumoModel",
	train      	 = train,
	test      	 = test,
	output_dir   = "./out/",
	max_epochs   = 1,
	batch_size   = 32, # On RTX 2080 Ti
	model = ViTForImageClassification.from_pretrained(
	    huggingface_model,
	    num_labels = len(label2id),
	    label2id   = label2id,
	    id2label   = id2label
	),
	feature_extractor = ViTFeatureExtractor.from_pretrained(
		huggingface_model,
	),
)

ref, hyp = trainer.evaluate_f1_score()

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

path = "./out/MyPneumoModel/20_2021-08-20-01-46-44/model/"
img  = "../../../samples/pneumothorax/with.jpg"

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

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