Hi, I followed a tutorial on how to train a dataset for Image Classification:
[link removed]
My problem is a usage example…
I amended the following, [link removed] to get the following code:
from PIL import Image
import numpy as np
import time
from torch import nn
from torch.nn import functional as F
#import glob
import os
import timm
# GEN LABELS
train_dir = "/home/ai/datasets/kaggle.com/ten_monkey_species/dataset_10_monkey_species/training/training/"
unique_labels = os.listdir(train_dir)
print(unique_labels)
class MonkeySpeciesClassifier(nn.Module):
def __init__(self, num_classes):
#define necessary layers
super().__init__()
self.num_classes = num_classes
self.model = timm.create_model(model_name = "resnet34", pretrained = True)
self.model.fc = nn.Linear(self.model.fc.in_features, out_features = num_classes)
def forward(self,X):
#define forward pass here
return F.softmax(self.model(X), dim=-1)
#model = MonkeySpeciesClassifier(len(unique_labels)).to(DEVICE)
#print(model(torch.zeros((1, 3, 256, 256)).to(DEVICE)).shape)
model = MonkeySpeciesClassifier(len(unique_labels))
print(model(torch.zeros((1, 3, 256, 256))).shape)
# Function to load the model
def load_model(model_path):
#model = torch.load(model_path)
#model = torch.load(model_path, encoding='ascii')
model.load_state_dict(torch.load(model_path))
if model:
model.eval() # Set the model to evaluation mode
else:
print("Error loading model: ", model_path)
return model
# Function to preprocess the image before feeding it to the model
def preprocess_image(image_path):
img = Image.open(image_path).convert("RGB")
img = img.resize((224, 224)) # Resize to the model's expected input size
img = np.array(img).astype(np.float32) / 255.0 # Convert to a NumPy array and normalize
img = np.transpose(img, (2, 0, 1)) # Transpose the image to (channels, height, width)
img = np.expand_dims(img, axis=0) # Add a batch dimension
return torch.tensor(img)
# Load the model
model_path = "best.pt"
model = load_model(model_path)
print("Loading image")
# Load the image and preprocess it
image_path = "/home/ai/datasets/kaggle.com/ten_monkey_species/dataset_10_monkey_species/validation/validation/n4/n404.jpg"
print("Processing image")
input_image = preprocess_image(image_path)
# Run the model
output = model(input_image)
print(output)
I’m not sure if I’m getting the wrong answer or no answer tbh… but the result is as follows, I tried to decode as I’ve seen, but as I say I’m still a noob, so any help / examples or even links to training and then using would be great please
tensor([[2.6537e-05, 6.2218e-06, 8.1168e-05, 7.7195e-05, 9.9924e-01, 1.3025e-04,
1.2812e-04, 4.9623e-05, 2.0674e-04, 4.9618e-05]],
grad_fn=<SoftmaxBackward0>)
Thanks