pupilsense / feature_extraction /features_extractor.py
vijul.shah
Input Video and Predictions as output video added
9acc552
raw
history blame
1.74 kB
import os
import sys
import warnings
import os.path as osp
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
sys.path.append(root_path)
from feature_extraction.extractor_mediapipe import ExtractorMediaPipe
warnings.filterwarnings("ignore")
class FeaturesExtractor:
def __init__(self, extraction_library="mediapipe", blink_detection=False, upscale=1):
self.upscale = upscale
self.blink_detection = blink_detection
self.extraction_library = extraction_library
self.feature_extractor = ExtractorMediaPipe(self.upscale)
def __call__(self, image):
results = {}
face = self.feature_extractor.extract_face(image)
if face is None:
# print("No face found. Skipped feature extraction!")
return None
else:
results["img"] = image
results["face"] = face
eyes_data = self.feature_extractor.extract_eyes(image, self.blink_detection)
if eyes_data is None:
# print("No eyes found. Skipped feature extraction!")
return results
else:
results["eyes"] = eyes_data
if eyes_data["blinked"]:
# print("Found blinked eyes!")
return results
else:
iris_data = self.feature_extractor.extract_iris(image)
if iris_data is None:
# print("No iris found. Skipped feature extraction!")
return results
else:
results["iris"] = iris_data
return results