| import torch | |
| import torch.nn as nn | |
| from torchvision import datasets, transforms, models | |
| import os | |
| from PIL import Image | |
| # Ustawienia parametrów modelu | |
| img_width, img_height = 224, 224 # Wymiary obrazu wymagane przez model ResNet | |
| model_path = './result/animal_classifier_resnet.pth' # Ścieżka do wytrenowanego modelu | |
| # Sprawdzenie, czy jest dostępny GPU i przypisanie urządzenia do zmiennej `device` | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}") | |
| # Wyłączenie CuDNN, chyba że twoja karta wspiera bibliotekę CuDNN (NVIDIA CUDA Deep Neural Network library) | |
| # torch.backends.cudnn.enabled = False | |
| # Transformacje danych wejściowych (zmiana rozmiaru, normalizacja) | |
| transform = transforms.Compose([ | |
| transforms.Resize((img_width, img_height)), # Zmiana rozmiaru obrazu | |
| transforms.ToTensor(), # Konwersja obrazu do tensoru | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # Normalizacja obrazu | |
| ]) | |
| # Inicjalizacja modelu ResNet18 bez wstępnych wag | |
| model = models.resnet18(weights=None) | |
| num_ftrs = model.fc.in_features # Liczba wejściowych cech ostatniej warstwy | |
| model.fc = nn.Linear(num_ftrs, len(datasets.ImageFolder('raw-img', transform=transform).classes)) # Zastąpienie ostatniej warstwy dopasowanej do liczby klas w danych | |
| # Przeniesienie modelu na GPU, jeśli jest dostępny | |
| model = model.to(device) | |
| model.load_state_dict(torch.load(model_path, map_location=device)) # Wczytanie wytrenowanych wag modelu | |
| model.eval() # Ustawienie modelu w tryb ewaluacyjny | |
| # Funkcja do przewidywania klasy obrazu | |
| def predict(image_path): | |
| image = Image.open(image_path).convert('RGB') # Otworzenie obrazu i konwersja do RGB | |
| image = transform(image).unsqueeze(0) # Zastosowanie transformacji i dodanie wymiaru batch | |
| image = image.to(device) # Przeniesienie obrazu na GPU, jeśli jest dostępny | |
| with torch.no_grad(): # Wyłączenie gradientów dla przewidywania | |
| output = model(image) # Przekazanie obrazu przez model | |
| _, predicted = torch.max(output, 1) # Wybranie klasy z najwyższym prawdopodobieństwem | |
| return datasets.ImageFolder('raw-img', transform=transform).classes[predicted.item()] # Zwrócenie nazwy klasy | |
| # Przykład użycia funkcji predict do przewidywania klasy obrazów w katalogu `recognize` | |
| test_image_dir = 'recognize' # Ścieżka do katalogu z obrazami do przewidywania | |
| for filename in os.listdir(test_image_dir): # Pętla przez pliki w katalogu | |
| if filename.lower().endswith(('.jpg', '.jpeg', '.png')): # Sprawdzenie rozszerzenia pliku | |
| image_path = os.path.join(test_image_dir, filename) # Pełna ścieżka do pliku | |
| print(f"Prediction for {filename}: {predict(image_path)}") # Wyświetlenie przewidywanej klasy dla obrazu | |