preprocessing/dataset_creation.py

import sys
import cv2
import os.path as osp

root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
sys.path.append(root_path)

from preprocessing.dataset_creation_utils import get_sr_method
from feature_extraction.features_extractor import FeaturesExtractor


class EyeDentityDatasetCreation:

    def __init__(self, feature_extraction_configs, sr_configs=None):
        self.extraction_library = feature_extraction_configs["extraction_library"]
        self.sr_configs = sr_configs
        if self.sr_configs:
            self.sr_method_name = sr_configs["method"]
            self.upscale = sr_configs["params"]["upscale"]
            if self.sr_method_name != "-":
                self.sr_method = get_sr_method(self, sr_configs)
        else:
            self.upscale = 1

        self.blink_detection = feature_extraction_configs["blink_detection"]
        self.features_extractor = FeaturesExtractor(
            extraction_library=self.extraction_library,
            blink_detection=self.blink_detection,
            upscale=self.upscale,
        )

    def __call__(self, img):
        # img = cv2.imread(img)
        if self.sr_configs is None or self.sr_configs != "-":
            img = img
        else:
            img = self.sr_method(img)

        result_dict = self.features_extractor(img)
        return result_dict