Spaces:
Sleeping
Sleeping
| 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() | |