garmentiq.landmark.derivation.prepare_args

 1"""Assembling the arguments a derivation function needs."""
 2import numpy as np
 3from .derivation_dict import derivation_dict
 4
 5
 6def prepare_args(
 7    entry: dict, derivation_dict: dict = derivation_dict, **extra_args
 8) -> dict:
 9    """
10    Prepares arguments for a specific derivation function based on an entry from the
11    landmark derivation configuration and extra arguments.
12
13    This function ensures that all required arguments for a given derivation function
14    are collected and validated against its schema defined in `derivation_dict`.
15    It also adds any necessary `extra_args` (like `landmark_coords` or `np_mask`).
16
17    Args:
18        entry (dict): A dictionary representing a single landmark's derivation entry,
19                      which must include a 'function' key and its specific parameters.
20        derivation_dict (dict): The global dictionary defining schemas for all derivation functions.
21                                Defaults to `derivation_dict`.
22        **extra_args: Additional keyword arguments that might be required by the derivation function,
23                      such as `landmark_coords` (NumPy array of landmark coordinates)
24                      and `np_mask` (NumPy array of the segmentation mask).
25
26    Raises:
27        ValueError: If the 'function' key is missing in `entry`, or if the function name is unknown.
28        TypeError: If an unsupported schema format is encountered for a key.
29        ValueError: If a required `extra_arg` is missing for a specific function (e.g., `landmark_coords`).
30
31    Returns:
32        dict: A dictionary where the key is the function name and its value is a dictionary
33              of arguments ready to be passed to that derivation function.
34    """
35    function_name = entry.get("function")
36    if function_name is None:
37        raise ValueError("Entry must include a 'function' key.")
38    if function_name not in derivation_dict:
39        raise ValueError(f"Unknown function: {function_name}")
40
41    function_schema = derivation_dict[function_name]
42    args = {}
43
44    for key, value in entry.items():
45        if key == "function":
46            continue
47
48        expected_type = function_schema.get(key)
49
50        if expected_type is None:
51            # Should be cast to int
52            args[key] = int(value)
53        elif isinstance(expected_type, list):
54            # Should be one of the listed options
55            if value not in expected_type:
56                raise ValueError(
57                    f"Invalid value '{value}' for {key}; expected one of {expected_type}"
58                )
59            args[key] = value
60        else:
61            raise TypeError(
62                f"Unsupported schema format for key '{key}' in function '{function_name}'"
63            )
64
65    # Add function-specific extra arguments
66    if function_name == "derive_keypoint_coord":
67        if "landmark_coords" not in extra_args:
68            raise ValueError(
69                "'landmark_coords' is required for 'derive_keypoint_coord'"
70            )
71        elif not isinstance(extra_args["landmark_coords"], np.ndarray):
72            raise ValueError("'landmark_coords' must be a 'np.ndarray'")
73
74        if "np_mask" not in extra_args:
75            raise ValueError("'np_mask' is required for 'derive_keypoint_coord'")
76        elif not isinstance(extra_args["np_mask"], np.ndarray):
77            raise ValueError("'np_mask' must be a 'np.ndarray'")
78        args["landmark_coords"] = extra_args["landmark_coords"]
79        args["np_mask"] = extra_args["np_mask"]
80        return {"derive_keypoint_coord": args}
81    # Add more if conditions if there are more derivation functions in the future
82    # elif function_name == "another_function_1":
83    #   if "arg_3" not in extra_args:
84    #     raise ValueError("'arg_3' is required for 'another_function_1'")
85    # else:
86    #     args['mask_path'] = extra_args['mask_path']
87    #     return args