The Five Artists

from datasets import load_dataset

dataset = load_dataset("eleon360/five-artists-dataset")
test_dataset = load_dataset("eleon360/five-artists-test-dataset")

dataset

test_dataset

# prompt: show me a 10x10 image with the first 50 images of the dataset together with their corresponding artist

import matplotlib.pyplot as plt
import numpy as np

# Assuming the dataset structure is consistent, access the 'train' split
train_dataset = dataset['train']

# Get the first 50 samples
first_100_samples = train_dataset.select(range(50))

# Create a 10x5 grid
fig, axes = plt.subplots(5, 10, figsize=(20, 10))
axes = axes.ravel() # Flatten the 2D array of axes

for i, ax in enumerate(axes):
    if i < len(first_100_samples):
        sample = first_100_samples[i]
        image = sample['image']
        artist = f"Artist {sample['artist_id']}"

        ax.imshow(image)
        ax.set_title(artist, fontsize=8)
        ax.axis('off')
    else:
        ax.axis('off') # Turn off empty subplots if less than 100 samples

plt.tight_layout()
plt.show()

# prompt: show us how to analyse the image data of the previous dataset.
# Import the necessary libraries
import matplotlib.pyplot as plt
import numpy as np

# Access the first image of the training set to analyse it on its own
# `train_dataset` was already loaded in the previous code
sample_image = train_dataset[0]

print("Analysing the first image")

# Extract the image and the artist ID from the sample
image = sample_image['image']
artist_id = sample_image['artist_id']

# Convert the image to a NumPy array so that its properties can be analysed
# The downloaded image is a PIL object; converting it to NumPy lets us treat it as numerical data
image_array = np.array(image)

# Print the shape of the image array (height, width, colour channels)
# This tells us the dimensions of the image in pixels
print(f"Image dimensions (height, width, colour channels): {image_array.shape}")

# Print the data type of the pixels
# They are usually 8-bit unsigned integers (uint8), which represent values from 0 to 255
print(f"Pixel data type: {image_array.dtype}")

# Compute and show the minimum and maximum pixel values in the image
# This gives us an idea of the range of colour intensities
print(f"Minimum pixel value: {image_array.min()}")
print(f"Maximum pixel value: {image_array.max()}")

# We can compute the mean pixel value to get a general idea of the brightness of the image
print(f"Mean pixel value: {image_array.mean():.2f}")

# We can also analyse the colour channels separately (if the image is in colour)
if image_array.shape[-1] == 3: # Check whether there are 3 channels (RGB)
    print("\nAnalysis by colour channel (Red, Green, Blue):")
    for i, color in enumerate(["Red", "Green", "Blue"]):
        # Compute the mean intensity of each colour channel
        print(f"  Mean of the {color} channel: {image_array[:,:,i].mean():.2f}")

# prompt: train a simple neural network to identify the artist number of each image.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import functional as F

# Define the desired size for resizing the images
IMG_HEIGHT = 32
IMG_WIDTH = 32

# Define a custom dataset to preprocess the images
class ArtistDataset(Dataset):
    def __init__(self, hf_dataset):
        self.hf_dataset = hf_dataset

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        item = self.hf_dataset[idx]
        # Convert the PIL image to a tensor and resize it
        image = item['image']
        image_tensor = F.to_tensor(image)
        image_resized = F.resize(image_tensor, (IMG_HEIGHT, IMG_WIDTH))

        # Normalise the pixel values (optional but recommended)
        image_normalized = image_resized * 255.0 # Scale from 0 to 255
        # Apply a simple normalisation (you can adjust this)
        mean = image_normalized.mean()
        std = image_normalized.std()
        image_normalized = (image_normalized - mean) / (std + 1e-6) # Avoids division by zero

        label = item['artist_id']
        # Make sure the label is a tensor of type long
        label_tensor = torch.tensor(label, dtype=torch.long)

        # Flatten the image
        image_flattened = image_normalized.view(-1)


        return image_flattened, label_tensor

# Create instances of your custom dataset for training and evaluation
train_dataset_custom = ArtistDataset(dataset['train'])


# Create DataLoaders to iterate over the data
# Use a small batch size so that it runs faster on CPU
BATCH_SIZE = 64
train_dataloader = DataLoader(train_dataset_custom, batch_size=BATCH_SIZE, shuffle=True)


# Define the super simple neural network (dense/linear layers only)
class SimpleNN(nn.Module):
    def __init__(self, input_size, num_classes):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, 128) # Dense layer
        self.relu = nn.ReLU()             # ReLU activation function
        self.fc2 = nn.Linear(128, 64)       # Another dense layer
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(64, num_classes) # Output layer (number of classes)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu2(x)
        x = self.fc3(x)
        return x

# Initialise the neural network
# The input size is IMG_HEIGHT * IMG_WIDTH * 3 (RGB channels)
INPUT_SIZE = IMG_HEIGHT * IMG_WIDTH * 3
NUM_CLASSES = 5 # There are 5 artists
model = SimpleNN(INPUT_SIZE, NUM_CLASSES)

# Define the loss function and the optimiser
criterion = nn.CrossEntropyLoss() # Loss for multiclass classification
optimizer = optim.Adam(model.parameters(), lr=0.01) # Adam optimiser

# Train the neural network
NUM_EPOCHS = 5 #

print("Starting training...")
for epoch in range(NUM_EPOCHS):
    model.train() # Put the model in training mode
    running_loss = 0.0
    correct_predictions = 0
    total_samples = 0

    for i, (inputs, labels) in enumerate(train_dataloader):
        optimizer.zero_grad() # Set the gradients to zero

        outputs = model(inputs)       # Pass the data through the network
        loss = criterion(outputs, labels) # Compute the loss

        loss.backward()       # Backpropagation
        optimizer.step()      # Update the weights

        running_loss += loss.item()

        # Compute the accuracy on the batch
        _, predicted = torch.max(outputs.data, 1)
        total_samples += labels.size(0)
        correct_predictions += (predicted == labels).sum().item()

        # Print progress from time to time
        if (i+1) % 100 == 0: # Print every 100 batches
            print(f'Epoch [{epoch+1}/{NUM_EPOCHS}], Step [{i+1}/{len(train_dataloader)}], Loss: {loss.item():.4f}')

    epoch_loss = running_loss / len(train_dataloader)
    epoch_accuracy = correct_predictions / total_samples
    print(f'Epoch [{epoch+1}/{NUM_EPOCHS}] Finished, Loss: {epoch_loss:.4f}, Train Accuracy: {epoch_accuracy:.4f}')

print("Training finished.")

# prompt: Generate predictions using the previous model for each of the images in test_dataset['train']

# Define a custom dataset for the test set
class ArtistTestDataset(Dataset):
    def __init__(self, hf_dataset):
        self.hf_dataset = hf_dataset

    def __len__(self):
        return len(self.hf_dataset['train']) # Access the 'train' split

    def __getitem__(self, idx):
        item = self.hf_dataset['train'][idx] # Access the 'train' split
        # Convert the PIL image to a tensor and resize it
        image = item['image']
        image_tensor = F.to_tensor(image)
        image_resized = F.resize(image_tensor, (IMG_HEIGHT, IMG_WIDTH))

        # Normalise the pixel values using the same statistics as in training (ideally)
        # To keep things simple, here we use the same simple normalisation
        image_normalized = image_resized * 255.0
        mean = image_normalized.mean()
        std = image_normalized.std()
        image_normalized = (image_normalized - mean) / (std + 1e-6)

        # Flatten the image
        image_flattened = image_normalized.view(-1)

        # There are no labels in the test set, so we return only the image
        return image_flattened

# Create an instance of the dataset
test_dataset_custom = ArtistTestDataset(test_dataset)

# Create a DataLoader
# We do not need shuffle for inference
test_dataloader = DataLoader(test_dataset_custom, batch_size=BATCH_SIZE, shuffle=False)

# Make the predictions
model.eval() # Put the model in evaluation mode (disables dropout, batchnorm, etc.)
predictions = []

#Generating predictions...
with torch.no_grad(): # We do not compute gradients during inference
    for inputs in test_dataloader:
        outputs = model(inputs)
        # Get the predicted class (the index with the highest probability)
        _, predicted = torch.max(outputs.data, 1)
        predictions.extend(predicted.tolist()) # Store the predictions as a list

print("Predictions generated.")

# You can print the first predictions to check them
print("First 10 predictions:", predictions[:10])

# prompt: generate image_ids with the image_ids in test_dataset

image_ids = [item['image_id'] for item in test_dataset['train']]
print("First 10 image_ids:", image_ids[:10])
print(f"Total number of image_ids: {len(image_ids)}")

# Create the submission.csv file
import csv

submission_filename = "submission.csv"

with open(submission_filename, mode='w', newline='') as file:
    writer = csv.writer(file)
    # Write the header
    writer.writerow(["image_id", "predicted_artist"])

    # Write the predictions
    for i in range(len(image_ids)):
        predicted_artist_id = predictions[i]
        # Map the predicted artist ID back to the artist name
        #predicted_artist_name = artist_id_to_name.get(predicted_artist_id, "Unknown") # Handle potential unknown IDs
        writer.writerow([int(image_ids[i]), predicted_artist_id])

print(f"Predictions saved to {submission_filename}")

# You can download the file from Colab if needed
# from google.colab import files
# files.download(submission_filename)