Introduction

Hunters and conservationists alike delight in understanding the health and wellness of the wildlife around them. Since capturing animals in the wild can prove difficult, tracking wellness of species and individuals is typically performed by photographing and diagnosing images of the wildlife. From the images, the age, stature, and wellness of different species can be discerned by wildlife experts.

Understanding the data

Ironically, there are thousands of trail camera images taken by hunters across the United States each year, but very few of them are seen by deer aging professionals. Of the images that are seen, even fewer see the light of day online or in publications. To alleviate the data drought, images used in this project are taken from a wide variety of sources including the National Deer Association (NDA), Field & Stream (F&S), state agencies, and other conservation resources.

In each case, the image of a male whitetail deer (“buck”) is captured along with the professional estimate of the deer’s age. Age and other information is stored in the image’s metadata, and ultimately used to create the truth labels within the convolution neural network (CNN) model. Although using images from a single source (ex. NDA) would ensure consistency across the panel of experts, utilization of multiple sources quickly buys us more data and allows for a wider swath of professionals to weigh in on the variables within an image that indicate the deer’s age.

Aspects of Aging

Interestingly, deer ages are always estimated in half years (ex. 1.5 years, 2.5 years, etc.) due to the life cycle of deer in hunting. Deer are typically born in the spring and harvested in the fall (deer hunting season), roughly half-way through their current age year. Newborns (fawns) are easy to identify — they are small, awkward looking, and have spots on their sides for the first 3-4 months of their lives. Since many hunters are interested in aging deer beyond the newborn stage, many images will feature deer at 1.5 year or older. Furthermore, the average lifetime of a deer in the wild ranges between three to six years. Taken together, we expect the labels for our supervised learning to lie within the set of [1.5, 2.5, 3.5, 4.5, 5.5] years, a very small discrete set of “correct” answers.

Along with each age ruling, the NDA provides a bit of justification; although not directly applicable to the ML model, their insights are helpful in building intuition into which features or patterns the ML model may eventually learn. Much like humans, each deer’s body grows and changes with age. For comparison, I’ve included a general comparison of features that are seen on young bucks versus mature males.

Understanding the data

Image Standardization

Trail cameras are often equipped with two viewing modes, “daytime” and “nighttime”. The daytime mode captures colored images containing red, green, and blue (RGB) pixel values, while the nighttime mode captures infrared grayscale images in grayscale. Furthermore, trail cameras are not created equal — they differ in their aspect ratios, pixel resolution, memory, motion sensitivity, and other features.

Before we expose our machine learning (ML) algorithm to the dataset, we want to make sure the dataset is as consistent as possible. Specifically, we want to remove variation between images (aspect ratio, coloration, etc.) within our dataset, so we begin by cropping out additional information in each picture. For landscape imagery, this means removing data tags at the bottom (the same information is stored in the filename anyways) and cropping dead space out of the image. If necessary, the still-rectangular image is combined with a white square to make the final image square, as shown in Figure 1.

Figure 1. Pre-conditioning the images

Once shaped, each image is stored with a five-part filename characterizing the image. Each image is assigned a filename of the format XXXXXX_ZZZZZZ_SS_NpN_P, where XXXXXX and ZZZZZZ denote the date the image was collected and the date the image was taken. Both dates are of the format YYMMDD (ex. March 31, 2025 would be represented as 250331). SS denotes the state the image was taken in (ex. Kentucky would be represented as KY), and NpN stands for the age of the deer (ex. 3p5 indicates an age of 3.5 years). Lastly, P represents the provider name (ex. “RLT” for Realtree, “NDA” for National Deer Association, etc.). Our images are now ready to be ingested.

Image Ingestion / formatting

We use glob to identify portable networks graphics (PNG) files within the image database, and we remove image files with no associated age (ex. files ending in NpN); this can occur when an image is gathered from one source (ex. NDA), but has not yet had its age revealed.

from glob import glob
from generic.analysis.basics import extract_labels
from generic.analysis.basics import ingest_resize_stack

# Find/ingest files in folder; force square & b/w
files = glob("..\\images\\squared\\*.png")
files = [s for s in files if "xpx" not in s]
print(len(files), "images found")

# Ingest images
images = ingest_resize_stack(files)
_,_,_,ages,_ = extract_labels(files)
print('Sample size:',images.shape)

# Convert labels to integers from 0 to 5 for proper one-hot encoding
# Create a mapping from your floating-point labels to integers
label_mapping = {label: i for i, label in enumerate(np.unique(ages))}
print("Label mapping:", label_mapping)

# Apply the mapping to convert labels to integers
integer_labels = np.array([label_mapping[l] for l in ages])
print("Converted labels:", integer_labels)

The identified valid files are ingested and stacked in a 3D array of (N images, X pixels, Y pixels) via ingest_resize_stack. The age of the deer in each photo is extracted from each image’s filename, and the ages (all in half years) are mapped to whole integers in two steps: (1) map float values to int values, then (2) apply the mapping to our deer ages.

To get an initial look at our dataset, images.shape displays the size of the data array. At the time of writing, this returns:

40 images found
Sample size: (40, 288, 288),

verifying that we have 40 individual images, each image reshaped and interpolated to a 288 x 288 two-dimensional grayscale image.

Train / test split

With the dataset loaded, we have yet to define our three key datasets: training data (X_train), test data (X_test), and validation data (X_valid). The first of these is the dataset we will use to train our ML model through forward- and back-propagation. Once the model is trained, we feed it validation data to see how the model reacts to data that was not in the original test dataset. Together, these datasets give us an idea of whether or model is under or over fit, allowing us to adjust our model selection and hyperparameters without biasing the final model. The test set is then used to provide an unbiased estimate of how well the final, tuned model will handle unseen data.

Following convention, we use an 80/20 split when separating the test data from the training data, and again when separating the validation data from the remaining training data. The grayscale images of X_train and X_test are then normalized between 0 and 255 to ensure each image shares the same range of values.

import keras
import numpy as np
from sklearn.model_selection import train_test_split

# Use a regular train_test_split without stratification
X_train, X_test, y_train, y_test = train_test_split(
    images, integer_labels, test_size=0.2, random_state=42
)

# Normalize the images
X_train = X_train.astype("float32") / 255.0
X_test = X_test.astype("float32") / 255.0

# Check the distribution after splitting
print("\nTraining set label distribution:")
unique_train_labels = np.unique(y_train)
for label in unique_train_labels:
    count = np.sum(y_train == label)
    print(f"Label {label} ({list(label_mapping.keys())[list(label_mapping.values()).index(label)]}): {count} samples ({count/len(y_train)*100:.2f}%)")

print("\nTest set label distribution:")
unique_test_labels = np.unique(y_test)
for label in unique_test_labels:
    count = np.sum(y_test == label)
    print(f"Label {label} ({list(label_mapping.keys())[list(label_mapping.values()).index(label)]}): {count} samples ({count/len(y_test)*100:.2f}%)")

# Get the number of classes
num_classes = len(label_mapping)

With our data and our labels in hand, we can then count how many instances of each age reside within our training and test datasets. According to these outputs, the age of 1.5 years is mapped to the integer 0, 2.5 years is mapped to integer 1, etc. There are two things to keep in mind. First, Python begins indexing at 0, which is why the first map element (1.5 years) maps to an index of 0. Second, this is not redefining “1.5 years” as “0 years” — instead, it is mapping the quantity “1.5 years” to some arbitrary numerical index we call “0”. We could easily just as map 1.5 years to the letter “A”, 2.5 years to the letter “B”, and it would achieve the same goal.

Label mapping: {np.float64(1.5): 0, np.float64(2.5): 1, np.float64(3.5): 2, ...}
Converted labels: [1 1 2 1 ... ]

Training set label distribution:
Label 0 (1.5): 3 samples (9.38%)
Label 1 (2.5): 9 samples (28.12%)
Label 2 (3.5): 8 samples (25.00%)
Label 3 (4.5): 3 samples (9.38%)
Label 4 (5.5): 7 samples (21.88%)
Label 5 (6.5): 1 samples (3.12%)
Label 6 (12.5): 1 samples (3.12%)

Although results are not shown for the test data, the second section shows the frequency of each label in the training data. At first glance, the values appear to be multi-moded — over half of the training samples contain deer aged at 2.5 or 3.5 years, while another 22% illustrate a deer aged 5.5 years.

As an aside, this is a side effect of how wildlife experts age deer. Although whitetail deer have been known to live as long as 22 years, many deer experts simply list a buck as “mature” once the deer reaches and exceeds an age of 5.5 years. For instance, a 7.5 year old buck will likely be judged by experts to be aged as “5.5 years”, “mature”, or “old”, even though his body continues to change over time.

Encode it like it’s hot

# Get the number of classes
num_classes = len(label_mapping)

# One-hot encode labels BEFORE splitting into train/validation
y_train_onehot = keras.utils.to_categorical(y_train, num_classes)
y_test_onehot = keras.utils.to_categorical(y_test, num_classes)

# Reshape data to add channel dimension
X_train = X_train.reshape(X_train.shape[0], 288, 288, 1)
X_test = X_test.reshape(X_test.shape[0], 288, 288, 1)

# Create a validation set (without stratification)
X_train_orig, X_valid, y_train_orig, y_valid = train_test_split(
    X_train, y_train_onehot, test_size=0.2, random_state=42
)

Migrating from discrete age values (1.5, 2.5, 3.5, …) to indices (0, 1, 2, …) shifts our modeling perspective; instead of predicting an age (ex. 2.5 years), we are instead predicting which category or “bucket” a picture should be dropped in. For instance, if an image contains a 2.5 year old buck, it should go in the second bucket; likewise, an image containing a 1.5 year old buck should go in the first bucket, and so on.

But there’s another step that needs to be performed to enable categorical prediction, and that is encoding, also referred to as “one hot” encoding. This coding method simply breaks up a list of N categories into N separate columns, and assigns each column a binary answer. To illustrate, consider a dataset of four images (below, left), where each image contains one animal — the first image contains a cat, the second contains a dog, the third contains another cat, and the fourth image contains a squirrel. After applying one hot encoding, our “Animal” column is replaced by three columns, one for each animal. The first image would have be assigned a “1” under “Cat” because it is an image of a cat, and “0” under the “Dog” and “Squirrel” columns, because the first image contains neither a dog nor a squirrel. The same logic continues for the other three images.

Figure 2. An illustration of one hot encoding

The same encoding is carried out on our dataset, except our data contains ages instead of animals. But why? This is a common technique to transform our categorical data (ex. “1.5 years”) into numerical format the machine learning algorithm can efficiently handle. Processing a pattern of 0’s and 1’s is much easier for pattern comparison than matching decimal values.

Data Augmentation

Our dataset is meager compared to datasets commonly used to train ML models. Typically, a dataset will include tens thousands of datapoints in stead of our 40; as we’ve discussed, this is due to the problem of availability. To alleviate this problem, we can use a technique called “data augmentation”, which does exactly what the name suggests — generate more data based off our original dataset.

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Print original sizes
print("\nBefore augmentation:")
print(X_train_orig.shape[0], "train samples")
print(X_test.shape[0], "test samples")
print(X_valid.shape[0], "validation samples")

# Setup data augmentation
datagen = ImageDataGenerator(
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True,
    vertical_flip=False
)

# Method 2: Generate augmented data in advance
# Define how many augmented samples per original sample
augmentation_factor = 5  # This will multiply your dataset by this factor
num_to_generate = X_train_orig.shape[0] * augmentation_factor

# Initialize empty arrays for augmented data
augmented_images = []
augmented_labels = []

# Create augmented images batch by batch
batch_size = 32
generated_count = 0

# Create a flow from the original data (without shuffling)
aug_gen = datagen.flow(
    X_train_orig, 
    y_train_orig,
    batch_size=batch_size,
    shuffle=False  # Important: keep the same order as labels
)

while generated_count < num_to_generate:
    # Get the next batch
    x_batch, y_batch = next(aug_gen)
    
    # Add to our collections
    augmented_images.append(x_batch)
    augmented_labels.append(y_batch)
    
    # Update the count
    generated_count += len(x_batch)
    
    # Break if we've generated enough
    if generated_count >= num_to_generate:
        break

# Concatenate all batches
augmented_images = np.concatenate(augmented_images)
augmented_labels = np.concatenate(augmented_labels)

# Trim excess (due to batch size)
augmented_images = augmented_images[:num_to_generate]
augmented_labels = augmented_labels[:num_to_generate]

# Combine with original data
X_train_combined = np.concatenate([X_train_orig, augmented_images])
y_train_combined = np.concatenate([y_train_orig, augmented_labels])

# Print new sizes after augmentation
print("\nAfter augmentation:")
print("Original training samples:", X_train_orig.shape[0])
print("Augmented training samples:", augmented_images.shape[0])
print("Combined training samples:", X_train_combined.shape[0])
print("Augmentation multiplier:", X_train_combined.shape[0] / X_train_orig.shape[0])
print("X_train_combined shape:", X_train_combined.shape)
print("y_train_combined shape:", y_train_combined.shape)

from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout

# Define the model
model = Sequential()
model.add(Conv2D(filters=16, kernel_size=3, padding='same', activation='relu', input_shape=(288, 288, 1)))
model.add(MaxPooling2D(pool_size=2))
model.add(Conv2D(filters=32, kernel_size=2, padding='same', activation='relu'))
model.add(MaxPooling2D(pool_size=2))
model.add(Conv2D(filters=64, kernel_size=2, padding='same', activation='relu'))
model.add(MaxPooling2D(pool_size=2))
model.add(Dropout(0.3))
model.add(Flatten())
model.add(Dense(500, activation='relu'))
model.add(Dropout(0.4))
model.add(Dense(num_classes, activation='softmax'))
model.summary()

model.compile(loss='categorical_crossentropy', optimizer='rmsprop', metrics=['accuracy'])

from keras.callbacks import ModelCheckpoint

# Train with pre-generated augmented data
checkpointer = ModelCheckpoint(filepath='model_augmented.weights.best.hdf5.keras', 
                             verbose=1, save_best_only=True)

print("\nTraining with pre-generated augmented data:")
hist_augmented = model.fit(
    X_train_combined, 
    y_train_combined,
    batch_size=32,
    epochs=20,
    validation_data=(X_valid, y_valid),
    callbacks=[checkpointer],
    verbose=1,
    shuffle=True
)

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix

# Evaluate with the best weights
model.load_weights('model_augmented.weights.best.hdf5.keras')

# Create a reverse mapping to get original labels
reverse_mapping = {i: label for label, i in label_mapping.items()}

# Make predictions
y_pred_prob = model.predict(X_test, verbose=0)
y_pred = np.argmax(y_pred_prob, axis=1)
y_true = np.argmax(y_test_onehot, axis=1)

# Create a comprehensive evaluation table that includes all classes
all_classes = list(range(num_classes))
all_class_names = [reverse_mapping[i] for i in all_classes]

# Create a comprehensive DataFrame
results_df = pd.DataFrame({
    'Class Index': all_classes,
    'Original Label': all_class_names,
    'In Test Set': [i in y_true for i in all_classes],
    'In Predictions': [i in y_pred for i in all_classes]
})

# Add metrics where applicable
precision_values = []
recall_values = []
f1_values = []
for cls in all_classes:
    if cls in y_true and cls in y_pred:
        # We can calculate metrics for this class
        true_binary = (y_true == cls).astype(int)
        pred_binary = (y_pred == cls).astype(int)
        precision_values.append(precision_score(true_binary, pred_binary, zero_division=0))
        recall_values.append(recall_score(true_binary, pred_binary, zero_division=0))
        f1_values.append(f1_score(true_binary, pred_binary, zero_division=0))
    else:
        # Class not present in test set or predictions
        precision_values.append(float('nan'))
        recall_values.append(float('nan'))
        f1_values.append(float('nan'))

results_df['Precision'] = precision_values
results_df['Recall'] = recall_values
results_df['F1 Score'] = f1_values

print("Comprehensive Class Evaluation:")
print(results_df.to_string(index=False))

# Create a confusion matrix (will only show classes present in test set)
cm = confusion_matrix(y_true, y_pred)
present_classes = sorted(set(np.concatenate([y_true, y_pred])))
plt.figure(figsize=(5, 4))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
           xticklabels=[f"{reverse_mapping[i]}" for i in present_classes],
           yticklabels=[f"{reverse_mapping[i]}" for i in present_classes])
plt.xlabel('Predicted')
plt.ylabel('True')
plt.title('Confusion Matrix (Only Classes Present in Test Set)')
plt.show()

# Show overall accuracy
accuracy = np.mean(y_pred == y_true)
print(f"\nOverall Test Accuracy: {accuracy:.4f}")

import keras
import numpy as np
from sklearn.model_selection import train_test_split

# Convert labels to integers from 0 to 5 for proper one-hot encoding
# Create a mapping from your floating-point labels to integers
label_mapping = {label: i for i, label in enumerate(np.unique(ages))}
print("Label mapping:", label_mapping)

# Apply the mapping to convert labels to integers
integer_labels = np.array([label_mapping[l] for l in ages])
print("Converted labels:", integer_labels)

# Use a regular train_test_split without stratification
X_train, X_test, y_train, y_test = train_test_split(
    images, integer_labels, test_size=0.2, random_state=42
)

# Check the distribution after splitting
print("\nTraining set label distribution:")
unique_train_labels = np.unique(y_train)
for label in unique_train_labels:
    count = np.sum(y_train == label)
    print(f"Label {label} ({list(label_mapping.keys())[list(label_mapping.values()).index(label)]}): {count} samples ({count/len(y_train)*100:.2f}%)")

print("\nTest set label distribution:")
unique_test_labels = np.unique(y_test)
for label in unique_test_labels:
    count = np.sum(y_test == label)
    print(f"Label {label} ({list(label_mapping.keys())[list(label_mapping.values()).index(label)]}): {count} samples ({count/len(y_test)*100:.2f}%)")

# Normalize the images
X_train = X_train.astype("float32") / 255.0
X_test = X_test.astype("float32") / 255.0

Model: "sequential"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                         ┃ Output Shape                ┃         Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
│ conv2d (Conv2D)                      │ (None, 288, 288, 16)        │             160 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d (MaxPooling2D)         │ (None, 144, 144, 16)        │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ conv2d_1 (Conv2D)                    │ (None, 144, 144, 32)        │           2,080 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d_1 (MaxPooling2D)       │ (None, 72, 72, 32)          │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ conv2d_2 (Conv2D)                    │ (None, 72, 72, 64)          │           8,256 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d_2 (MaxPooling2D)       │ (None, 36, 36, 64)          │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dropout (Dropout)                    │ (None, 36, 36, 64)          │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ flatten (Flatten)                    │ (None, 82944)               │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dense (Dense)                        │ (None, 500)                 │      41,472,500 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dropout_1 (Dropout)                  │ (None, 500)                 │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dense_1 (Dense)                      │ (None, 7)                   │           3,507 │
└──────────────────────────────────────┴─────────────────────────────┴─────────────────┘

 Total params: 41,486,503 (158.26 MB)

 Trainable params: 41,486,503 (158.26 MB)

 Non-trainable params: 0 (0.00 B)

Training with pre-generated augmented data:
Epoch 1/20
5/5 ━━━━━━━━━━━━━━━━━━━━ 0s 351ms/step - accuracy: 0.1925 - loss: 1.9045
Epoch 1: val_loss improved from inf to 1.97243, saving model to model_augmented.weights.best.hdf5.keras
5/5 ━━━━━━━━━━━━━━━━━━━━ 8s 2s/step - accuracy: 0.1916 - loss: 1.8951 - val_accuracy: 0.1429 - val_loss: 1.9724
Epoch 2/20
5/5 ━━━━━━━━━━━━━━━━━━━━ 0s 329ms/step - accuracy: 0.3223 - loss: 1.7746
Epoch 2: val_loss improved from 1.97243 to 1.94029, saving model to model_augmented.weights.best.hdf5.keras

Comprehensive Class Evaluation:
 Class Index  Original Label  In Test Set  In Predictions  Precision  Recall  F1 Score
           0             1.5         True           False        NaN     NaN       NaN
           1             2.5         True            True      0.375     1.0  0.545455
           2             3.5         True           False        NaN     NaN       NaN
           3             4.5         True           False        NaN     NaN       NaN
           4             5.5         True           False        NaN     NaN       NaN
           5             6.5        False           False        NaN     NaN       NaN
           6            12.5        False           False        NaN     NaN       NaN

Overall Test Accuracy: 0.3750

# Get the number of classes
num_classes = len(label_mapping)

# One-hot encode labels BEFORE splitting into train/validation
y_train_onehot = keras.utils.to_categorical(y_train, num_classes)
y_test_onehot = keras.utils.to_categorical(y_test, num_classes)

# Reshape data to add channel dimension
X_train = X_train.reshape(X_train.shape[0], 288, 288, 1)
X_test = X_test.reshape(X_test.shape[0], 288, 288, 1)

# Create a validation set (without stratification)
X_train_orig, X_valid, y_train_orig, y_valid = train_test_split(
    X_train, y_train_onehot, test_size=0.2, random_state=42
)

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Print original sizes
print("\nBefore augmentation:")
print(X_train_orig.shape[0], "train samples")
print(X_test.shape[0], "test samples")
print(X_valid.shape[0], "validation samples")

# Setup data augmentation
datagen = ImageDataGenerator(
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True,
    vertical_flip=False
)

# Method 2: Generate augmented data in advance
# Define how many augmented samples per original sample
augmentation_factor = 5  # This will multiply your dataset by this factor
num_to_generate = X_train_orig.shape[0] * augmentation_factor

# Initialize empty arrays for augmented data
augmented_images = []
augmented_labels = []

# Create augmented images batch by batch
batch_size = 32
generated_count = 0

# Create a flow from the original data (without shuffling)
aug_gen = datagen.flow(
    X_train_orig, 
    y_train_orig,
    batch_size=batch_size,
    shuffle=False  # Important: keep the same order as labels
)

while generated_count < num_to_generate:
    # Get the next batch
    x_batch, y_batch = next(aug_gen)
    
    # Add to our collections
    augmented_images.append(x_batch)
    augmented_labels.append(y_batch)
    
    # Update the count
    generated_count += len(x_batch)
    
    # Break if we've generated enough
    if generated_count >= num_to_generate:
        break

# Concatenate all batches
augmented_images = np.concatenate(augmented_images)
augmented_labels = np.concatenate(augmented_labels)

# Trim excess (due to batch size)
augmented_images = augmented_images[:num_to_generate]
augmented_labels = augmented_labels[:num_to_generate]

# Combine with original data
X_train_combined = np.concatenate([X_train_orig, augmented_images])
y_train_combined = np.concatenate([y_train_orig, augmented_labels])

# Print new sizes after augmentation
print("\nAfter augmentation:")
print("Original training samples:", X_train_orig.shape[0])
print("Augmented training samples:", augmented_images.shape[0])
print("Combined training samples:", X_train_combined.shape[0])
print("Augmentation multiplier:", X_train_combined.shape[0] / X_train_orig.shape[0])
print("X_train_combined shape:", X_train_combined.shape)
print("y_train_combined shape:", y_train_combined.shape)

import keras
import numpy as np
from sklearn.model_selection import train_test_split

# Convert labels to integers from 0 to 5 for proper one-hot encoding
# Create a mapping from your floating-point labels to integers
label_mapping = {label: i for i, label in enumerate(np.unique(labels))}
print("Label mapping:", label_mapping)

# Apply the mapping to convert labels to integers
integer_labels = np.array([label_mapping[l] for l in labels])
print("Converted labels:", integer_labels)

# Use a regular train_test_split without stratification
X_train, X_test, y_train, y_test = train_test_split(
    images, integer_labels, test_size=0.2, random_state=42
)

# Check the distribution after splitting
print("\nTraining set label distribution:")
unique_train_labels = np.unique(y_train)
for label in unique_train_labels:
    count = np.sum(y_train == label)
    print(f"Label {label} ({list(label_mapping.keys())[list(label_mapping.values()).index(label)]}): {count} samples ({count/len(y_train)*100:.2f}%)")

print("\nTest set label distribution:")
unique_test_labels = np.unique(y_test)
for label in unique_test_labels:
    count = np.sum(y_test == label)
    print(f"Label {label} ({list(label_mapping.keys())[list(label_mapping.values()).index(label)]}): {count} samples ({count/len(y_test)*100:.2f}%)")

# Normalize the images
X_train = X_train.astype("float32") / 255.0
X_test = X_test.astype("float32") / 255.0

# One-hot encode labels with the correct number of classes
num_classes = len(label_mapping)
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# Create a validation set (without stratification)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_train, y_train, test_size=0.2, random_state=42
)

# Reshape data to add channel dimension
X_train = X_train.reshape(X_train.shape[0], 288, 288, 1)
X_valid = X_valid.reshape(X_valid.shape[0], 288, 288, 1)
X_test = X_test.reshape(X_test.shape[0], 288, 288, 1)

print(X_train.shape[0], "train samples")
print(X_test.shape[0], "test samples")
print(X_valid.shape[0], "validation samples")

from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout

# Define the model
model = Sequential()
model.add(Conv2D(filters=16, kernel_size=3, padding='same', activation='relu', input_shape=(288, 288, 1)))
model.add(MaxPooling2D(pool_size=2))
model.add(Conv2D(filters=32, kernel_size=2, padding='same', activation='relu'))
model.add(MaxPooling2D(pool_size=2))
model.add(Conv2D(filters=64, kernel_size=2, padding='same', activation='relu'))
model.add(MaxPooling2D(pool_size=2))
model.add(Dropout(0.3))
model.add(Flatten())
model.add(Dense(500, activation='relu'))
model.add(Dropout(0.4))
model.add(Dense(num_classes, activation='softmax'))
model.summary()

model.compile(loss='categorical_crossentropy', optimizer='rmsprop', metrics=['accuracy'])

model.compile(loss='categorical_crossentropy', optimizer='rmsprop', metrics=['accuracy'])

from keras.callbacks import ModelCheckpoint

checkpointer = ModelCheckpoint(filepath='model.weights.best.hdf5.keras', verbose=1, save_best_only=True)
hist = model.fit(X_train, y_train, batch_size=5, epochs=100, validation_data=(X_valid,y_valid), callbacks=[checkpointer], verbose=4, shuffle=True)

from keras.callbacks import ModelCheckpoint

checkpointer = ModelCheckpoint(filepath='model.weights.best.hdf5.keras', verbose=1, save_best_only=True)
hist = model.fit(X_train, y_train, batch_size=5, epochs=100, validation_data=(X_valid,y_valid), callbacks=[checkpointer], verbose=4, shuffle=True)