import os
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from PIL import Image
from models import ResNet18
from datasets import HandGestureDataset

# Set the path to the dataset directory
data_dir = 'dataset'

# Define the image transforms
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Define the dataset
dataset = HandGestureDataset(data_dir, transform=transform)

# Define the data loader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# Load the pre-trained neural network
model = ResNet18(pretrained=True)

# Replace the final fully connected layer with a new one
num_classes = 7
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

# Set the model to evaluation mode
model.eval()

# Load an image from the dataset
for i, (image, label) in enumerate(dataloader):
    # Apply the image transforms
    image = transform(image)

    # Add a batch dimension
    image = image.unsqueeze(0)

    # Make a prediction on the image
    with torch.no_grad():
        output = model(image)
        prediction = torch.argmax(output)

    # Print the prediction
    print(f'Image {i+1}, Predicted note: {prediction.item()}')