train_knn.py

import os
import numpy as np
from PIL import Image  # ✅ Replaces cv2
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import joblib  # For saving model
from sklearn.utils.multiclass import unique_labels

# Helper function to get all image paths recursively
def getListOfFiles(dirName):
    listOfFile = os.listdir(dirName)
    allFiles = []
    for entry in listOfFile:
        fullPath = os.path.join(dirName, entry)
        if os.path.isdir(fullPath):
            allFiles += getListOfFiles(fullPath)
        else:
            allFiles.append(fullPath)
    return allFiles

# 📁 Path to your custom dataset (change if needed)
DATASET_PATH = "./knn_number_dataset/"
imagePaths = getListOfFiles(DATASET_PATH)

# 📊 Load and preprocess images
data = []
labels = []
for i, imagePath in enumerate(imagePaths):
    label = os.path.basename(os.path.dirname(imagePath))  # Folder name is the label

    try:
        # ✅ Using PIL instead of cv2
        with Image.open(imagePath) as img:
            img = img.convert("RGB")  # Ensure 3 channels
            img = img.resize((32, 32))  # Resize to 32x32
            image = np.array(img)
    except Exception as e:
        print(f"⚠️ Skipping unreadable image: {imagePath} | Error: {e}")
        continue

    data.append(image)
    labels.append(label)

    if i % 100 == 0:
        print(f"📥 Processed {i} images...")

# 🔢 Convert to NumPy arrays
data = np.array(data)
labels = np.array(labels)

# 🧠 Encode string labels
le = LabelEncoder()
labels = le.fit_transform(labels)

# 🧼 Flatten image data
dataset_size = data.shape[0]
data = data.reshape(dataset_size, -1)

print(f"✅ Data shape: {data.shape}")
print(f"✅ Labels shape: {labels.shape}")
print(f"✅ Unique classes: {set(labels)}")

# 🔀 Train/test split
(trainX, testX, trainY, testY) = train_test_split(
    data, labels, test_size=0.25, random_state=42, stratify=labels
)

# 🧠 Train KNN model
model = KNeighborsClassifier(n_neighbors=3, n_jobs=-1)
model.fit(trainX, trainY)

# 📊 Evaluate
print("\n📋 Classification Report:")
predictions = model.predict(testX)
present_labels = unique_labels(testY, predictions)

print(classification_report(
    testY,
    predictions,
    labels=present_labels,
    target_names=le.inverse_transform(present_labels)
))

# 💾 Save model and label classes
os.makedirs("Models", exist_ok=True)
joblib.dump(model, os.path.join("Models", "knn_model.pkl"))
np.save(os.path.join("Models", "label_classes.npy"), le.classes_)

print("✅ Model saved to: Models/knn_model.pkl")
print("✅ Label classes saved to: Models/label_classes.npy")