Spaces:
Sleeping
Sleeping
File size: 4,909 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 | import cv2
import mediapipe as mp
import numpy as np
import os
import time
from sign_utils import MediaPipeLandmarksModel
class DataCollector:
def __init__(self):
self.landmark_model = MediaPipeLandmarksModel()
self.mp_drawing = mp.solutions.drawing_utils
self.mp_hands = mp.solutions.hands
def collect_sign_data(self, sign_name, num_sequences=30, sequence_length=30):
"""Collect data for a specific sign."""
# Create directory for the sign
data_dir = f"dataset/{sign_name}"
os.makedirs(data_dir, exist_ok=True)
cap = cv2.VideoCapture(0)
# Arrays to store sequences
sequences = []
print(f"Collecting data for sign: {sign_name}")
print(f"Will collect {num_sequences} sequences of {sequence_length} frames each")
print("Press 's' to start collecting, 'q' to quit")
collecting = False
sequence = []
sequence_count = 0
frame_count = 0
while True:
ret, frame = cap.read()
if not ret:
break
# Flip frame horizontally for mirror effect
frame = cv2.flip(frame, 1)
# Extract landmarks
landmarks = self.landmark_model.predict_landmarks(frame)
# Draw landmarks
if landmarks:
frame = self.landmark_model.draw_landmarks(frame, landmarks)
# Display status
if collecting:
cv2.putText(frame, f'Collecting: {sign_name}', (20, 50),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
cv2.putText(frame, f'Sequence: {sequence_count + 1}/{num_sequences}', (20, 100),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
cv2.putText(frame, f'Frame: {frame_count}/{sequence_length}', (20, 150),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
if landmarks:
# Extract features
features = self.landmark_model.extract_features(landmarks)
sequence.append(features)
frame_count += 1
if frame_count >= sequence_length:
# Save sequence
sequences.append(sequence.copy())
sequence = []
frame_count = 0
sequence_count += 1
if sequence_count >= num_sequences:
break
# Wait before next sequence
time.sleep(2)
else:
cv2.putText(frame, f'Ready to collect: {sign_name}', (20, 50),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
cv2.putText(frame, "Press 's' to start", (20, 100),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
cv2.imshow('Data Collection', frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('s') and not collecting:
collecting = True
print("Started collecting...")
elif key == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
# Save collected data
if sequences:
np.save(f"{data_dir}/sequences.npy", np.array(sequences))
print(f"Saved {len(sequences)} sequences for {sign_name}")
return sequences
def main():
collector = DataCollector()
# List of signs to collect
signs = ['hello', 'thank_you', 'please', 'sorry', 'yes', 'no', 'good', 'bad', 'help', 'stop']
print("Sign Language Data Collection Tool")
print("Available signs:", signs)
while True:
print("\nOptions:")
print("1. Collect data for a specific sign")
print("2. Collect data for all signs")
print("3. Exit")
choice = input("Enter your choice (1-3): ")
if choice == '1':
sign_name = input("Enter sign name: ").lower().replace(' ', '_')
num_sequences = int(input("Number of sequences (default 30): ") or 30)
collector.collect_sign_data(sign_name, num_sequences)
elif choice == '2':
for sign in signs:
print(f"\nPrepare to collect data for: {sign}")
input("Press Enter when ready...")
collector.collect_sign_data(sign, 30)
elif choice == '3':
break
else:
print("Invalid choice!")
if __name__ == "__main__":
main()
|