import torch
import urllib.request
from PIL import Image
from datasets import load_dataset
from transformers import EfficientNetImageProcessor, EfficientNetForImageClassification

dataset = load_dataset("chriamue/bird-species-dataset")

#####
labels = dataset["test"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = str(i)
    id2label[str(i)] = label

preprocessor = EfficientNetImageProcessor.from_pretrained("google/efficientnet-b2")
model = EfficientNetForImageClassification.from_pretrained("chriamue/bird-species-classifier", num_labels=len(
    labels), id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True)

image = dataset["validation"][0]["image"]

url = 'https://upload.wikimedia.org/wikipedia/commons/a/a9/Common_Blackbird.jpg'
image = Image.open(urllib.request.urlretrieve(url)[0])


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

inputs = preprocessor(image, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}

with torch.no_grad():
    logits = model(**inputs).logits
    predicted_label = logits.argmax(-1).item()
    print(labels[predicted_label])