# -*- coding: utf-8 -*-
"""Final Run-Concurrent

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/19foLOwCXRH0e0P_Xqqc-9VgpnmjYyAX8
"""

# Install the Necessary Packages

!pip install datasets huggingface_hub sentence-transformers gradio evaluate
!pip install git+https://github.com/huggingface/accelerate
!pip install transformers==4.28.0

import datasets
from datasets import load_dataset
import pandas
from PIL import Image
import cv2
import os
from pandas import read_csv
from google.colab import drive

drive.mount('/content/drive/')

raw_dataset = load_dataset("imagefolder", data_dir="/content/drive/MyDrive/california_fire_damage_classification_merged/train")
dataset = raw_dataset["train"].train_test_split(test_size=0.2, stratify_by_column="label")

from transformers import ViTImageProcessor, ViTForImageClassification
import torch

device = 'cuda' # for GPU

model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

model.eval()
#model.to(device);

# image_processor is the same as Tokenizer
extractor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')

labels = raw_dataset['train'].features['label'].names

labels

from transformers import AutoFeatureExtractor, AutoModelForImageClassification, AutoTokenizer

extractor = AutoFeatureExtractor.from_pretrained("/content/drive/MyDrive/california_fire_damage_classification_merged/saved_model_files")
model = AutoModelForImageClassification.from_pretrained("/content/drive/MyDrive/california_fire_damage_classification_merged/saved_model_files")

import torch

def transform(example_batch):
    inputs = extractor([x.convert("RGB") for x in example_batch['image']], return_tensors='pt')
    inputs['labels'] = example_batch['label']
    return inputs

prepared_ds = dataset.with_transform(transform)

### RUNNING EVALUATION ON PRETRAINED MODEL

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments("test_trainer"),


import numpy as np
from datasets import load_metric

metric = load_metric("accuracy")


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)


trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=None,
    eval_dataset=prepared_ds['test'],
    compute_metrics=compute_metrics,
)

j = 2095

print('Groundtruth: ', y_test_np[j], ' ', labels[y_test_np[j]], 'Prediction: ', y_predicts_np[j], ' ', labels[y_predicts_np[j]])
dataset['test'][j]['image']

pixel_values_array = []
y_test = []
counter = 0

for img_pair in prepared_ds['test']:
  pixel_values_array.append(img_pair['pixel_values'])
  y_test.append(img_pair["labels"])
  #pixel_values_tensor = torch.concat((pixel_values_tensor, img_pair['pixel_values']), 0)
  counter += 1
  print(counter)

#pixel_values_tensor = torch.stack(pixel_values_array)

#pixel_values_tensor

len(pixel_values_tensor)
len(y_predicts_merged)

import numpy as np
y_test_np = np.array(y_test)
y_predicts_np = np.array(y_predicts_merged)

np.where((y_test_np == y_predicts_np) == False)

y_predicts = []

for i in range(len(pixel_values_tensor)):
  logits = model(pixel_values_tensor[i:i+1])[-1]
  y_predict = [logit.argmax(-1).item() for logit in logits]
  y_predicts.append(y_predict)

y_predicts

y_predicts_merged = [inner for outer in y_predicts for inner in outer]

y_predicts_merged

logits = model(pixel_values_tensor[0:1])[-1]

logits

y_predict = [logit.argmax(-1).item() for logit in logits]
y_predict

#y_test = [img_pair["labels"] for img_pair in prepared_ds['test']]
y_test = prepared_ds['test'][0:100]["labels"]
y_test

from sklearn.metrics import classification_report, confusion_matrix

print(confusion_matrix(y_test, y_predicts_merged))
print(classification_report(y_test, y_predicts_merged))

probability = torch.nn.functional.softmax(logits, dim=-1)
probability

probs = probability.detach().numpy()
probs

confidences = [{label: float(prob[j]) for j, label in enumerate(labels)} for prob in probs]
confidences

# First we get the features corresponding to the first training image
encoding = image_processor(images=prepared_ds['test'][0]['image'], return_tensors="pt").to(device)

# Then pass it through the model and get a prediction

######
outputs = model(**encoding)
logits = outputs.logits
######

prediction = logits.argmax(-1).item()

print("Predicted class:", model.config.id2label[prediction])

# For 1 Sample -> look at distribution of probabilities assigned

tokenizer = AutoTokenizer.from_pretrained("google/vit-base-patch16-224")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

encoding = image_processor(images=[prepared_ds["test"][0]['image']], return_tensors="pt").to(device)
outputs = model(**encoding)
logits = outputs.logits
prediction = logits.argmax(-1).item()

print("Predicted class:", model.config.id2label[prediction])

im_test = [dataset['test'][0]['image'], dataset['test'][1]['image']]
features_test = extractor(im_test, return_tensors='pt')
features_test['pixel_values'][0]



features_test['pixel_values'][-1]

logits = model(features_test["pixel_values"])
logits[-1]

probability = torch.nn.functional.softmax(logits, dim=-1)

logits = model(features_test["pixel_values"])[-1]

probs = probability[0].detach().numpy()
confidences = {label: float(probs[i]) for i, label in enumerate(labels)}

probability = torch.nn.functional.softmax(logits, dim=-1)
probability

prepared_ds['test'][0]['pixel_values']