72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import os.path
|
|
from collections import OrderedDict
|
|
from glob import glob
|
|
from typing import Generator, Union, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
HandMask = np.ndarray
|
|
HandMaskMaybeWithFileName = Union[HandMask, Tuple[HandMask, str]]
|
|
|
|
_CSV_FILE_NAME = 'data.csv'
|
|
_CSV_COLUMNS = OrderedDict([('frame_number', int),
|
|
('hand_id', int),
|
|
('x', float),
|
|
('y', float),
|
|
('direction_x', float),
|
|
('direction_y', float),
|
|
('distance', float),
|
|
('detected', bool),
|
|
('frames_without_detection', int)])
|
|
|
|
_HAND_MASK_FILE_GLOB = 'hand_labels_*.png'
|
|
|
|
|
|
def load_coordinates(data_csv_path: str) -> pd.DataFrame:
|
|
data_csv_path = os.path.abspath(data_csv_path)
|
|
hand_data = pd.read_csv(data_csv_path,
|
|
sep=',',
|
|
header=None,
|
|
names=_CSV_COLUMNS.keys(),
|
|
dtype=_CSV_COLUMNS)
|
|
|
|
return hand_data
|
|
|
|
|
|
def load_hand_masks(data_dir_path: str,
|
|
return_file_names: bool = False) -> Generator[HandMaskMaybeWithFileName, None, None]:
|
|
data_dir_path = os.path.abspath(data_dir_path)
|
|
|
|
mask_files = glob(os.path.join(data_dir_path, _HAND_MASK_FILE_GLOB))
|
|
mask_files = sorted(mask_files)
|
|
|
|
for mf in mask_files:
|
|
hand_mask = cv2.imread(mf, cv2.IMREAD_UNCHANGED)
|
|
|
|
if return_file_names:
|
|
yield hand_mask, mf
|
|
else:
|
|
yield hand_mask
|
|
|
|
|
|
class RecordingHandTracking:
|
|
def __init__(self, path: str):
|
|
self._path = os.path.abspath(path)
|
|
|
|
def hand_coordinates(self) -> pd.DataFrame:
|
|
return load_coordinates(os.path.join(self._path, _CSV_FILE_NAME))
|
|
|
|
def hand_mask_sequence(self, return_file_names: bool = False) -> Generator[HandMaskMaybeWithFileName, None, None]:
|
|
return load_hand_masks(self._path, return_file_names)
|
|
|
|
|
|
for mf in mask_files:
|
|
hand_mask = cv2.imread(mf, cv2.IMREAD_UNCHANGED)
|
|
|
|
if return_file_names:
|
|
yield hand_mask, mf
|
|
else:
|
|
yield hand_mask
|