Spaces:
Sleeping
Sleeping
File size: 5,179 Bytes
1f4e4ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | 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()
|