EasyTinyML / train.py
harisankarrj's picture
Upload folder using huggingface_hub
b0fdb22 verified
from tensorflow.keras import layers, models, optimizers
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras import layers, models, optimizers
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import os
# Function to load images and labels from folders
def load_data(folder_path):
images = []
labels = []
for label in os.listdir(folder_path):
label_path = os.path.join(folder_path, label)
if os.path.isdir(label_path):
for filename in os.listdir(label_path):
img_path = os.path.join(label_path, filename)
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # Convert to RGB
images.append(img)
labels.append(label)
return np.array(images), np.array(labels)
# Load data from folders
data_path = r"C://Users//Admin//Downloads//webappml//data"
images, labels = load_data(data_path)
# Encode labels
label_encoder = LabelEncoder()
encoded_labels = label_encoder.fit_transform(labels)
# Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(images, encoded_labels, test_size=0.2, random_state=42)
# Normalize pixel values to be between 0 and 1
X_train, X_test = X_train / 255.0, X_test / 255.0
# Data Augmentation
datagen = ImageDataGenerator(
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True,
fill_mode='nearest'
)
# Fit the ImageDataGenerator on the training data
datagen.fit(X_train)
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(len(set(labels)), activation='softmax'))
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# Train the model with augmented data
model.fit(datagen.flow(X_train, y_train, batch_size=32), epochs=10, validation_data=(X_test, y_test))
# Save the model
model.save("harisankar.h5")
print("Model trained with data augmentation and saved successfully.")