# Importowanie bibliotek
# Torch - głębokie sieci neuronowe (ang. Deep Neural Networks)
import torch
from torchvision import transforms
from PIL import Image
import timm
 
# Wczytaj model Swin Transformer z biblioteki timm
# ta wersja jest wyuczona na zbiorze ImageNet-21k
nazwa_modelu = "swin_large_patch4_window7_224"  
klasyfikator = timm.create_model(nazwa_modelu, pretrained=True)
# klasyfikator.eval()
 
# Pobierz zbiór ImageNet (1000 klas)
import json
imagenet_labels_url = "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json"
import requests
labels = json.loads(requests.get(imagenet_labels_url).text)
 
# Potok przetwarzania obrazu na potrzeby klasyfikatora
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),  # model Swin Transformer bierze na wejściu obrazy 224x224
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),  # Normalizacja
])
 
 
def recognize_image(img_path):
    """Wczytuje obraz z pliku, przetwarza zgodnie z potokiem preprocess a na koniec klasyfikuje"""
    img = Image.open(img_path).convert("RGB")
    # Dodatkowy wymiar na wejściu
    img_tensor = preprocess(img).unsqueeze(0)  
   
    # Predykcja, czyli przypisanie prawdopodobieństw klas do obrazu wejściowego
    with torch.no_grad():
        logits = klasyfikator(img_tensor)
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        top5_prob, top5_catid = torch.topk(probabilities, 5)
 
    # Przeliczenie prawdopodobieństw na etykiety klas
    results = [(labels[catid], prob.item()) for catid, prob in zip(top5_catid[0], top5_prob[0])]
    return results
 
 
if __name__ == "__main__":
    img_path = "obrazek.jpg"
    wyjscie_klasyfikatora = recognize_image(img_path)
    for i, (label, confidence) in enumerate(wyjscie_klasyfikatora):
        print(f"{i + 1}: {label} ({confidence:.2f})")