Skip to content

Segmentation IOProcessor Plugin#

This plugin targets segmentation tasks and allows for the input image to be split in tiles of a size that depends on the model and on an arbitrary number of bands.

During initialization, the plugin accesses the model's data module configuration from the vLLM configuration and instantiates a DataModule object dynamically.

This plugin is installed as terratorch_segmentation.

Plugin specification#

Model requirements#

This plugin expects the model to take two parameters for inference. The first, named pixel_values, points to a tensor containing the raw image data extracted from the input tiff. The second parameter, named location_coords, is optional and points to a tensor containing geospatial coordinates for the image.

Below an example input model specification accepted by this plugin. The user can change the shapes of the tensors according to their model requirements but the number and names of the fields must be kept unchanged.

Model input specification accepted by the Segmentation IOProcessor plugin
"input":{
    "target": "pixel_values",
    "data":{
        "pixel_values":{
            "type": "torch.Tensor",
            "shape": [6, 512, 512]
        },
        "location_coords":{
            "type":"torch.Tensor",
            "shape": [1, 2]
        }
    }
}

Full details on TerraTorch models input model specification for vLLM are available here.

Plugin configuration#

This plugin allows for additional configuration data to be passed via the TERRATORCH_SEGMENTATION_IO_PROCESSOR_CONFIG environment variable. If set, the variable should contain the plugin configuration in json string format.

The plugin configuration format is defined in the PluginConfig class.

terratorch.vllm.plugins.segmentation.types.PluginConfig #

Bases: BaseModel

Source code in terratorch/vllm/plugins/segmentation/types.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class PluginConfig(BaseModel):
    output_path: str = None
    """
    Default output folder path to be used when the out_data_format is set to path. 
    If omitted, the plugin will default to the current user home directory.
    """

    @model_validator(mode="after")
    def validate_values(self) -> Self:
        if not self.output_path:
            self.output_path = str(Path.home())
        elif os.path.exists(self.output_path):
            if not os.access(self.output_path, os.W_OK):
                raise ValueError(f"The path: {self.output_path} is not writable")
        else:
            raise ValueError(f"The path: {self.output_path} does not exist")

        return self

output_path = None class-attribute instance-attribute #

Default output folder path to be used when the out_data_format is set to path. If omitted, the plugin will default to the current user home directory.

Request Data Format#

The input format for the plugin is defined in the RequestData class.

terratorch.vllm.plugins.segmentation.types.RequestData #

Bases: BaseModel

Source code in terratorch/vllm/plugins/segmentation/types.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class RequestData(BaseModel):

    data_format: Literal["b64_json", "path", "url"]
    """
    Data type for the input image.
    Allowed values are: [`b64_json`, `path`, `url`]
    """

    out_data_format: Literal["b64_json", "path"]
    """
    Data type for the output image.
    Allowed values are: [`b64_json`, `path`]
    """

    data: Any
    """
    Input image data
    """

    indices: Optional[list[int]] = None
    """
    Indices for bands to be processed in the input file
    """

data instance-attribute #

Input image data

data_format instance-attribute #

Data type for the input image. Allowed values are: [b64_json, path, url]

indices = None class-attribute instance-attribute #

Indices for bands to be processed in the input file

out_data_format instance-attribute #

Data type for the output image. Allowed values are: [b64_json, path]

Depending on the values set in data_format, the plugin expects data to contain a string that complies to the format. Similarly, out_data_format controls the data format returned to the user. The field indices can be customised by the user and it is expected to be a list of integers.

Request Output Format#

The output format for the plugin is defined in the RequestOutput class.

terratorch.vllm.plugins.segmentation.types.RequestOutput #

Bases: BaseModel

Source code in terratorch/vllm/plugins/segmentation/types.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class RequestOutput(BaseModel):

    data_format: Literal["b64_json", "path"]
    """
    Data type for the output image.
    Allowed values are: [`b64_json`, `path`]
    """

    data: Any
    """
    Output image data
    """

    request_id: Optional[str] = None
    """
    The vLLM request ID if applicable
    """

data instance-attribute #

Output image data

data_format instance-attribute #

Data type for the output image. Allowed values are: [b64_json, path]

request_id = None class-attribute instance-attribute #

The vLLM request ID if applicable

Plugin Defaults#

Tiled Inference Parameters#

By default the plugin uses the same horizontal and vertical crop value of 512 when computing image tiles. Users can use different crop values by specifying them in their model config.json file. See the example below that overrides the default values with vertical and horizontal crop values of 256.

Custom tiled inference parameters in model configuration
{
  "pretrained_cfg": {
    "model": {
      "init_args": {
        "tiled_inference_parameters": {
          "h_crop": 256,
          "w_crop": 256
        }
      }
    }
  }
}

Please note, the tiled_inference_parameters field is not mandatory in the model configuration. Full details on the model configuration file can be found here.

Image Input Indices#

By default the plugin extracts bands at indices [0, 1, 2, 3, 4, 5] from the input image. The user can customise this for each image, by setting the indices field accordingly in the inference request payload.