Noob Image Classification Training Play

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

train_dir = "/home/ai/datasets/"
unique_labels = os.listdir(train_dir)

class MonkeySpeciesClassifier(nn.Module):
	def __init__(self, num_classes):
		#define necessary layers
		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')
	if model:
		model.eval()  # Set the model to evaluation mode
		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 ="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 = ""
model = load_model(model_path)

print("Loading image")
# Load the image and preprocess it
image_path = "/home/ai/datasets/"

print("Processing image")
input_image = preprocess_image(image_path)

# Run the model
output = model(input_image)

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]],


I’m unsure where exactly you are stuck, but in case you want to interpret the output of your model you would interpret it as a tensor containing probabilities for each class given the input image. In your case the model predicts with a high probability the sample belongs to class4. Let me know if I misunderstood your question.

Hi, sorry was a stressful night, but oh! oops!

Many thanks :laughing: