mlm/stac_model/input.py at main · Nazim-crim/mlm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
from typing import Annotated, Literal, TypeAlias, Union
from typing_extensions import Self

from pydantic import Field, model_validator

from stac_model.base import (
    DataType,
    MLMBaseModel,
    ModelBandsOrVariablesReferences,
    Number,
    OmitIfNone,
    ProcessingExpression,
)


class InputStructure(MLMBaseModel):
    shape: list[Union[int, float]] = Field(min_length=1)
    dim_order: list[str] = Field(min_length=1)
    data_type: DataType

    @model_validator(mode="after")
    def validate_dimensions(self) -> Self:
        if len(self.shape) != len(self.dim_order):
            raise ValueError("Dimension order and shape must be of equal length for corresponding indices.")
        return self


class ValueScalingClipMin(MLMBaseModel):
    type: Literal["clip-min"] = "clip-min"
    minimum: Number


class ValueScalingClipMax(MLMBaseModel):
    type: Literal["clip-max"] = "clip-max"
    maximum: Number


class ValueScalingClip(MLMBaseModel):
    type: Literal["clip"] = "clip"
    minimum: Number
    maximum: Number


class ValueScalingMinMax(MLMBaseModel):
    type: Literal["min-max"] = "min-max"
    minimum: Number
    maximum: Number


class ValueScalingZScore(MLMBaseModel):
    type: Literal["z-score"] = "z-score"
    mean: Number
    stddev: Number


class ValueScalingOffset(MLMBaseModel):
    type: Literal["offset"] = "offset"
    value: Number


class ValueScalingScale(MLMBaseModel):
    type: Literal["scale"] = "scale"
    value: Number


class ValueScalingProcessingExpression(ProcessingExpression):
    type: Literal["processing"] = "processing"


ValueScalingObject: TypeAlias = Union[
    ValueScalingMinMax,
    ValueScalingZScore,
    ValueScalingClip,
    ValueScalingClipMin,
    ValueScalingClipMax,
    ValueScalingOffset,
    ValueScalingScale,
    ValueScalingProcessingExpression,
    None,
]

ResizeType: TypeAlias = (
    Literal[
        "crop",
        "pad",
        "interpolation-nearest",
        "interpolation-linear",
        "interpolation-cubic",
        "interpolation-area",
        "interpolation-lanczos4",
        "interpolation-max",
        "wrap-fill-outliers",
        "wrap-inverse-map",
    ]
    | None
)


class ModelInput(ModelBandsOrVariablesReferences):
    name: str
    input: InputStructure
    value_scaling: Annotated[list[ValueScalingObject] | None, OmitIfNone] = None
    resize_type: Annotated[ResizeType | None, OmitIfNone] = None
    pre_processing_function: ProcessingExpression | list[ProcessingExpression] | None = None