AirHacking-Security / train_model.py
kuldeep0204's picture
Create train_model.py
33096a8 verified
# train_model.py
import numpy as np
import os
from glob import glob
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout, BatchNormalization
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
import joblib
DATA_DIR = "gesture_data"
SEQ_LEN = 30
MODEL_OUT = "gesture_lstm.h5"
LABELS_OUT = "labels.joblib"
def load_data(data_dir):
X = []
y = []
for label in sorted(os.listdir(data_dir)):
p = os.path.join(data_dir, label)
if not os.path.isdir(p): continue
files = glob(os.path.join(p, "*.npz"))
for f in files:
d = np.load(f)['data']
if d.shape[0] != SEQ_LEN:
continue
X.append(d) # shape (seq_len, features)
y.append(label)
X = np.array(X)
y = np.array(y)
return X, y
def build_model(input_shape, num_classes):
model = Sequential([
LSTM(128, return_sequences=True, input_shape=input_shape),
Dropout(0.3),
LSTM(64),
BatchNormalization(),
Dropout(0.3),
Dense(64, activation="relu"),
Dense(num_classes, activation="softmax")
])
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
return model
def main():
X, y = load_data(DATA_DIR)
print("Loaded X:", X.shape, "y:", y.shape)
le = LabelEncoder()
y_enc = le.fit_transform(y)
joblib.dump(le, LABELS_OUT)
X_train, X_test, y_train, y_test = train_test_split(X, y_enc, test_size=0.2, random_state=42, stratify=y_enc)
# normalize coords per sample (optional)
# Flatten features along coordinates scale? We'll standardize by sample mean/std
# Here we scale coords to zero mean unit std per sample for robustness
def normalize_samples(arr):
arr2 = arr.copy()
for i in range(arr2.shape[0]):
s = arr2[i]
mean = s.mean(axis=0)
std = s.std(axis=0) + 1e-8
arr2[i] = (s - mean) / std
return arr2
X_train = normalize_samples(X_train)
X_test = normalize_samples(X_test)
input_shape = (X_train.shape[1], X_train.shape[2]) # (seq_len, features)
model = build_model(input_shape, len(le.classes_))
print(model.summary())
ckpt = ModelCheckpoint(MODEL_OUT, save_best_only=True, monitor="val_loss", verbose=1)
es = EarlyStopping(monitor="val_loss", patience=10, restore_best_weights=True)
history = model.fit(X_train, y_train, validation_data=(X_test, y_test),
epochs=100, batch_size=16, callbacks=[ckpt, es])
# final evaluation
loss, acc = model.evaluate(X_test, y_test)
print("Test loss/acc:", loss, acc)
print("Saved model to", MODEL_OUT)
print("Saved labels to", LABELS_OUT)
if __name__ == "__main__":
main()