|
|
import cv2
|
|
|
import numpy as np
|
|
|
import csv
|
|
|
import sys
|
|
|
import time
|
|
|
from pathlib import Path
|
|
|
import argparse
|
|
|
import pandas as pd
|
|
|
|
|
|
#sys.path.append(str(Path(__file__).parent.parent))
|
|
|
from skeleton.mediapipe_detector import MediaPipeDetector
|
|
|
from ml_gestures_dynamic.feature_extractor import extract_sequence
|
|
|
|
|
|
def main():
|
|
|
parser = argparse.ArgumentParser()
|
|
|
parser.add_argument('--label', required=True, help='Название жеста (например, wave_left)')
|
|
|
parser.add_argument('--output', default='dynamic_data.csv', help='CSV файл для сохранения')
|
|
|
parser.add_argument('--camera', type=int, default=0)
|
|
|
parser.add_argument('--duration', type=float, default=3.0, help='Длительность записи (сек)')
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
output_path = Path(args.output)
|
|
|
|
|
|
# Создаём файл с заголовком, если его ещё нет
|
|
|
if not output_path.exists():
|
|
|
with open(output_path, 'w', newline='') as f:
|
|
|
writer = csv.writer(f)
|
|
|
header = ['label', 'sequence_id', 'frame_idx'] + [f'f{i}' for i in range(99)]
|
|
|
writer.writerow(header)
|
|
|
|
|
|
detector = MediaPipeDetector()
|
|
|
cap = cv2.VideoCapture(args.camera)
|
|
|
if not cap.isOpened():
|
|
|
print("Камера не найдена")
|
|
|
sys.exit(1)
|
|
|
|
|
|
# Определяем следующий ID последовательности
|
|
|
try:
|
|
|
df = pd.read_csv(output_path)
|
|
|
next_id = df['sequence_id'].max() + 1 if not df.empty else 0
|
|
|
except:
|
|
|
next_id = 0
|
|
|
|
|
|
print(f"Запись жеста: {args.label}. Нажмите SPACE для начала, q для выхода.")
|
|
|
recording = False
|
|
|
start_time = 0
|
|
|
sequence = []
|
|
|
|
|
|
while True:
|
|
|
ret, frame = cap.read()
|
|
|
if not ret:
|
|
|
break
|
|
|
frame = cv2.flip(frame, 1) # зеркало для удобства
|
|
|
result = detector.detect(frame)
|
|
|
if result['success']:
|
|
|
landmarks = result['landmarks']
|
|
|
vis = detector.draw_landmarks(frame, result['pose_landmarks'])
|
|
|
else:
|
|
|
vis = frame
|
|
|
landmarks = None
|
|
|
|
|
|
# Если мы в режиме записи
|
|
|
if recording:
|
|
|
if landmarks is not None:
|
|
|
sequence.append(landmarks)
|
|
|
elapsed = time.time() - start_time
|
|
|
# Отображаем прогресс записи
|
|
|
cv2.putText(vis, f"RECORDING... {elapsed:.1f}/{args.duration}", (10, 30),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,255), 2)
|
|
|
if elapsed >= args.duration:
|
|
|
recording = False
|
|
|
seq_features = extract_sequence(sequence)
|
|
|
with open(output_path, 'a', newline='') as f:
|
|
|
writer = csv.writer(f)
|
|
|
for i, feat in enumerate(seq_features):
|
|
|
writer.writerow([args.label, next_id, i] + feat.tolist())
|
|
|
print(f"Сохранено {len(sequence)} кадров для жеста {args.label}, ID={next_id}")
|
|
|
sequence = []
|
|
|
next_id += 1
|
|
|
else:
|
|
|
# Ожидание нажатия пробела
|
|
|
cv2.putText(vis, f"Press SPACE to record '{args.label}'", (10, 30),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,0), 2)
|
|
|
|
|
|
cv2.imshow('Record dynamic gesture', vis)
|
|
|
key = cv2.waitKey(1) & 0xFF
|
|
|
if key == ord('q'):
|
|
|
break
|
|
|
if key == ord(' ') and not recording:
|
|
|
# Обратный отсчёт на живом видео
|
|
|
for i in range(3, 0, -1):
|
|
|
# Получаем свежий кадр для отсчёта
|
|
|
ret, frame = cap.read()
|
|
|
if not ret:
|
|
|
break
|
|
|
frame = cv2.flip(frame, 1)
|
|
|
result = detector.detect(frame)
|
|
|
if result['success']:
|
|
|
vis = detector.draw_landmarks(frame, result['pose_landmarks'])
|
|
|
else:
|
|
|
vis = frame
|
|
|
cv2.putText(vis, f"Starting in {i}...", (frame.shape[1]//2-100, frame.shape[0]//2),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0,0,255), 3)
|
|
|
cv2.imshow('Record dynamic gesture', vis)
|
|
|
cv2.waitKey(1000)
|
|
|
# Показываем "GO!"
|
|
|
ret, frame = cap.read()
|
|
|
if ret:
|
|
|
frame = cv2.flip(frame, 1)
|
|
|
result = detector.detect(frame)
|
|
|
if result['success']:
|
|
|
vis = detector.draw_landmarks(frame, result['pose_landmarks'])
|
|
|
else:
|
|
|
vis = frame
|
|
|
cv2.putText(vis, "GO!", (frame.shape[1]//2-50, frame.shape[0]//2),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 2, (0,255,0), 3)
|
|
|
cv2.imshow('Record dynamic gesture', vis)
|
|
|
cv2.waitKey(500)
|
|
|
# Начинаем запись
|
|
|
recording = True
|
|
|
start_time = time.time()
|
|
|
sequence = []
|
|
|
|
|
|
cap.release()
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
main()
|