|
|
import cv2
|
|
|
import os
|
|
|
import csv
|
|
|
import sys
|
|
|
from pathlib import Path
|
|
|
import argparse
|
|
|
|
|
|
#sys.path.append(str(Path(__file__).parent.parent))
|
|
|
from skeleton.mediapipe_detector import MediaPipeDetector
|
|
|
from ml_gestures.feature_extractor import normalize_landmarks
|
|
|
|
|
|
def main():
|
|
|
parser = argparse.ArgumentParser(description='Разметка изображений для обучения')
|
|
|
parser.add_argument('--folder', required=True, help='Папка с изображениями')
|
|
|
parser.add_argument('--classes', default='dome,cross,none',
|
|
|
help='Список классов через запятую')
|
|
|
parser.add_argument('--output', default='gesture_data.csv',
|
|
|
help='Имя выходного CSV-файла')
|
|
|
parser.add_argument('--max_display_size', default='800,600',
|
|
|
help='Максимальный размер для отображения (ширина,высота)')
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
classes = [c.strip() for c in args.classes.split(',')]
|
|
|
key_to_class = {str(i+1): cls for i, cls in enumerate(classes)}
|
|
|
print("Классы:", classes)
|
|
|
|
|
|
detector = MediaPipeDetector()
|
|
|
|
|
|
max_width, max_height = map(int, args.max_display_size.split(','))
|
|
|
|
|
|
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
|
|
|
image_files = [f for f in os.listdir(args.folder) if f.lower().endswith(image_extensions)]
|
|
|
image_files.sort()
|
|
|
print(f"Найдено {len(image_files)} изображений.")
|
|
|
|
|
|
csv_file = args.output
|
|
|
file_exists = os.path.isfile(csv_file)
|
|
|
if not file_exists:
|
|
|
with open(csv_file, 'w', newline='', encoding='utf-8') as f:
|
|
|
writer = csv.writer(f)
|
|
|
writer.writerow(['class'] + [f'f{i}' for i in range(99)])
|
|
|
|
|
|
for idx, filename in enumerate(image_files):
|
|
|
filepath = os.path.join(args.folder, filename)
|
|
|
print(f"\n[{idx+1}/{len(image_files)}] {filename}")
|
|
|
|
|
|
image = cv2.imread(filepath)
|
|
|
if image is None:
|
|
|
print("Не удалось загрузить")
|
|
|
continue
|
|
|
|
|
|
result = detector.detect(image)
|
|
|
|
|
|
if result['success']:
|
|
|
landmarks = result['landmarks']
|
|
|
features = normalize_landmarks(landmarks)
|
|
|
vis_image = detector.draw_landmarks(image.copy(), result['pose_landmarks'])
|
|
|
else:
|
|
|
vis_image = image.copy()
|
|
|
landmarks = None
|
|
|
|
|
|
# Масштабирование для отображения
|
|
|
h, w = vis_image.shape[:2]
|
|
|
scale = min(max_width / w, max_height / h, 1.0)
|
|
|
if scale < 1.0:
|
|
|
new_w = int(w * scale)
|
|
|
new_h = int(h * scale)
|
|
|
display = cv2.resize(vis_image, (new_w, new_h))
|
|
|
else:
|
|
|
display = vis_image.copy()
|
|
|
|
|
|
# Панель с инструкцией
|
|
|
dh, dw = display.shape[:2]
|
|
|
overlay = display.copy()
|
|
|
cv2.rectangle(overlay, (0, dh-80), (dw, dh), (50,50,50), -1)
|
|
|
cv2.addWeighted(overlay, 0.6, display, 0.4, 0, display)
|
|
|
|
|
|
y = dh - 60
|
|
|
for i, cls in enumerate(classes):
|
|
|
cv2.putText(display, f"{i+1}:{cls}", (10 + i*120, y),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2)
|
|
|
cv2.putText(display, "n:skip q:quit", (10, y+30),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,0), 2)
|
|
|
|
|
|
cv2.imshow('Annotation', display)
|
|
|
key = cv2.waitKey(0) & 0xFF
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if key == ord('q'):
|
|
|
break
|
|
|
elif key == ord('n'):
|
|
|
continue
|
|
|
else:
|
|
|
key_char = chr(key) if key < 256 else None
|
|
|
if key_char in key_to_class:
|
|
|
selected = key_to_class[key_char]
|
|
|
with open(csv_file, 'a', newline='', encoding='utf-8') as f:
|
|
|
writer = csv.writer(f)
|
|
|
writer.writerow([selected] + features.tolist())
|
|
|
print(f"Сохранено: {selected}")
|
|
|
else:
|
|
|
print("Неверная клавиша")
|
|
|
|
|
|
print("Готово.")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
main()
|