Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import tensorflow as tf | |
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import Dense, Dropout, BatchNormalization | |
| import cv2 | |
| import os | |
| import glob | |
| from sign_utils import MediaPipeLandmarksModel | |
| def load_and_process_dataset(): | |
| """Load digit images and extract landmarks.""" | |
| landmark_model = MediaPipeLandmarksModel() | |
| X_landmarks = [] | |
| y = [] | |
| print("Loading and processing images...") | |
| # Process each digit folder | |
| for digit in range(10): | |
| digit_path = os.path.join("dataset", str(digit)) | |
| if not os.path.exists(digit_path): | |
| continue | |
| image_files = glob.glob(os.path.join(digit_path, "*.JPG")) | |
| print(f"Processing {len(image_files)} images for digit {digit}") | |
| count = 0 | |
| for img_file in image_files: | |
| # Load image | |
| img = cv2.imread(img_file) | |
| if img is None: | |
| continue | |
| # Extract landmarks | |
| landmarks = landmark_model.predict_landmarks(img) | |
| if landmarks is not None: | |
| # Normalize landmarks | |
| landmarks_normalized = landmark_model.normalize_landmarks(landmarks) | |
| X_landmarks.append(landmarks_normalized) | |
| y.append(digit) | |
| count += 1 | |
| # Limit to 100 images per digit for faster training | |
| if count >= 100: | |
| break | |
| print(f"Successfully processed {count} images for digit {digit}") | |
| print(f"Total processed: {len(X_landmarks)} images with landmarks") | |
| # Check if we have any data | |
| if len(X_landmarks) == 0: | |
| return np.array([]), np.array([]) | |
| # Ensure all landmark arrays have the same shape | |
| landmark_shape = X_landmarks[0].shape[0] | |
| filtered_landmarks = [] | |
| filtered_y = [] | |
| for i, landmarks in enumerate(X_landmarks): | |
| if landmarks.shape[0] == landmark_shape: | |
| filtered_landmarks.append(landmarks) | |
| filtered_y.append(y[i]) | |
| print(f"After filtering: {len(filtered_landmarks)} valid samples") | |
| return np.array(filtered_landmarks), np.array(filtered_y) | |
| def create_simple_model(num_classes=10): | |
| """Create a simple neural network for landmark classification.""" | |
| model = Sequential([ | |
| Dense(128, activation='relu', input_shape=(63,)), # 21 landmarks * 3 coordinates | |
| BatchNormalization(), | |
| Dropout(0.3), | |
| Dense(64, activation='relu'), | |
| BatchNormalization(), | |
| Dropout(0.3), | |
| Dense(32, activation='relu'), | |
| Dropout(0.2), | |
| Dense(num_classes, activation='softmax') | |
| ]) | |
| model.compile( | |
| optimizer='adam', | |
| loss='sparse_categorical_crossentropy', | |
| metrics=['accuracy'] | |
| ) | |
| return model | |
| def main(): | |
| print("Simple Sign Language Digits Training") | |
| print("=" * 40) | |
| # Load dataset | |
| X, y = load_and_process_dataset() | |
| if len(X) == 0: | |
| print("No valid images found with hand landmarks!") | |
| return | |
| # Split data (simple split without sklearn) | |
| n_samples = len(X) | |
| n_train = int(0.8 * n_samples) | |
| # Shuffle data | |
| indices = np.random.permutation(n_samples) | |
| train_indices = indices[:n_train] | |
| test_indices = indices[n_train:] | |
| X_train, X_test = X[train_indices], X[test_indices] | |
| y_train, y_test = y[train_indices], y[test_indices] | |
| print(f"Training set: {len(X_train)} samples") | |
| print(f"Test set: {len(X_test)} samples") | |
| # Create model | |
| model = create_simple_model() | |
| print("\nModel Summary:") | |
| model.summary() | |
| # Train model | |
| print("\nTraining model...") | |
| history = model.fit( | |
| X_train, y_train, | |
| validation_data=(X_test, y_test), | |
| epochs=30, | |
| batch_size=32, | |
| verbose=1 | |
| ) | |
| # Evaluate model | |
| print("\nEvaluating model...") | |
| test_loss, test_accuracy = model.evaluate(X_test, y_test, verbose=0) | |
| print(f"Test Accuracy: {test_accuracy:.4f}") | |
| # Save model | |
| print("\nSaving model...") | |
| # Create models directory if it doesn't exist | |
| os.makedirs('models', exist_ok=True) | |
| # Save Keras model | |
| model.save('models/digit_model.h5') | |
| # Convert to TFLite | |
| try: | |
| 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 successfully!") | |
| except Exception as e: | |
| print(f"TFLite conversion failed: {e}") | |
| # Save labels | |
| labels = [str(i) for i in range(10)] | |
| with open('models/labels.txt', 'w') as f: | |
| f.write('\n'.join(labels)) | |
| print("Training completed successfully!") | |
| print("You can now run the main application: streamlit run app.py") | |
| if __name__ == "__main__": | |
| main() | |