Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import tensorflow as tf | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.preprocessing import LabelEncoder | |
| import os | |
| import glob | |
| from utils.model_utils import create_model, get_training_callbacks, evaluate_model | |
| from utils.data_utils import DataProcessor | |
| class SignLanguageTrainer: | |
| def __init__(self): | |
| self.data_processor = DataProcessor() | |
| self.label_encoder = LabelEncoder() | |
| def load_dataset(self, dataset_path="dataset"): | |
| """Load dataset from collected data.""" | |
| X, y = [], [] | |
| labels = [] | |
| # Get all sign directories | |
| sign_dirs = [d for d in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, d))] | |
| for sign_name in sign_dirs: | |
| sign_path = os.path.join(dataset_path, sign_name) | |
| sequence_file = os.path.join(sign_path, "sequences.npy") | |
| if os.path.exists(sequence_file): | |
| sequences = np.load(sequence_file) | |
| for sequence in sequences: | |
| X.append(sequence) | |
| y.append(sign_name) | |
| if sign_name not in labels: | |
| labels.append(sign_name) | |
| print(f"Loaded {len(sequences)} sequences for {sign_name}") | |
| # Save labels | |
| with open('models/labels.txt', 'w') as f: | |
| f.write('\n'.join(labels)) | |
| # Encode labels | |
| y_encoded = self.label_encoder.fit_transform(y) | |
| return np.array(X), y_encoded, labels | |
| def prepare_data(self, X, y, test_size=0.2, val_size=0.2): | |
| """Prepare data for training.""" | |
| # Split into train and test | |
| X_train, X_test, y_train, y_test = train_test_split( | |
| X, y, test_size=test_size, random_state=42, stratify=y | |
| ) | |
| # Split train into train and validation | |
| X_train, X_val, y_train, y_val = train_test_split( | |
| X_train, y_train, test_size=val_size, random_state=42, stratify=y_train | |
| ) | |
| return X_train, X_val, X_test, y_train, y_val, y_test | |
| def train_model(self, X_train, y_train, X_val, y_val, | |
| num_classes, epochs=50, batch_size=32): | |
| """Train the sign language model.""" | |
| # Get input shape | |
| sequence_length, num_features = X_train.shape[1], X_train.shape[2] | |
| # Create model | |
| model = create_model(num_classes, sequence_length, num_features) | |
| print("Model Summary:") | |
| model.summary() | |
| # Get callbacks | |
| callbacks = get_training_callbacks() | |
| # Train model | |
| history = model.fit( | |
| X_train, y_train, | |
| validation_data=(X_val, y_val), | |
| epochs=epochs, | |
| batch_size=batch_size, | |
| callbacks=callbacks, | |
| verbose=1 | |
| ) | |
| return model, history | |
| def save_model(self, model, model_name="sign_language_model"): | |
| """Save the trained model.""" | |
| # Save Keras model | |
| model.save(f'models/{model_name}.h5') | |
| # Convert to TFLite | |
| converter = tf.lite.TFLiteConverter.from_keras_model(model) | |
| converter.optimizations = [tf.lite.Optimize.DEFAULT] | |
| # Quantization for smaller model size | |
| def representative_dataset(): | |
| for _ in range(100): | |
| yield [np.random.random((1, 30, 126)).astype(np.float32)] | |
| converter.representative_dataset = representative_dataset | |
| converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] | |
| converter.inference_input_type = tf.uint8 | |
| converter.inference_output_type = tf.uint8 | |
| try: | |
| tflite_model = converter.convert() | |
| with open('models/model.tflite', 'wb') as f: | |
| f.write(tflite_model) | |
| print("Model converted to TFLite successfully!") | |
| except Exception as e: | |
| print(f"TFLite conversion failed: {e}") | |
| # Fallback to float32 model | |
| converter = tf.lite.TFLiteConverter.from_keras_model(model) | |
| converter.optimizations = [tf.lite.Optimize.DEFAULT] | |
| tflite_model = converter.convert() | |
| with open('models/model.tflite', 'wb') as f: | |
| f.write(tflite_model) | |
| print("Model converted to TFLite (float32) successfully!") | |
| def main(): | |
| trainer = SignLanguageTrainer() | |
| print("Sign Language Model Training") | |
| print("=" * 40) | |
| # Check if dataset exists | |
| if not os.path.exists("dataset"): | |
| print("Dataset not found! Please run data_collection.py first.") | |
| return | |
| # Load dataset | |
| print("Loading dataset...") | |
| X, y, labels = trainer.load_dataset() | |
| if len(X) == 0: | |
| print("No data found! Please collect data first using data_collection.py") | |
| return | |
| print(f"Loaded {len(X)} sequences for {len(labels)} signs") | |
| print(f"Signs: {labels}") | |
| # Prepare data | |
| print("Preparing data...") | |
| X_train, X_val, X_test, y_train, y_val, y_test = trainer.prepare_data(X, y) | |
| print(f"Training set: {len(X_train)} sequences") | |
| print(f"Validation set: {len(X_val)} sequences") | |
| print(f"Test set: {len(X_test)} sequences") | |
| # Train model | |
| print("Training model...") | |
| model, history = trainer.train_model( | |
| X_train, y_train, X_val, y_val, | |
| num_classes=len(labels), | |
| epochs=50, | |
| batch_size=32 | |
| ) | |
| # Evaluate model | |
| print("Evaluating model...") | |
| results = evaluate_model(model, X_test, y_test) | |
| print(f"Test Accuracy: {results['accuracy']:.4f}") | |
| # Save model | |
| print("Saving model...") | |
| trainer.save_model(model) | |
| print("Training completed successfully!") | |
| print("You can now run the main application: streamlit run app.py") | |
| if __name__ == "__main__": | |
| main() | |