Standardize sample dict keys

[adamjstewart]

Summary

In TorchGeo, each Dataset returns a sample dictionary with key/value pairs for each object. I propose we standardize the names of these keys and the types and dimensions of their corresponding values as follows:

• input: Tensor, [C, H, W]
• label: Tensor, [1]
• mask: Tensor, [H, W]
• bbox: Tensor, [N, 4, 2], defined in clockwise order: top-left, top-right, bottom-right and bottom-left, coordinates in x, y order
• bbox_xyxy: Tensor, [N, 4], defined in xmin, ymin, xmax, ymax format
• bbox_xywh: Tensor, [N, 4], defined in xmin, ymin, width, height format
• keypoints: Tensor, [N, 2], coordinates in x, y order

I also propose we standardize batch dictionaries (mini-batches of samples) using the exact same definition, but with an additional batch dimension in front.

Rationale

Our current keys are not uniform (bbox and boxes, label and labels):

$ grep -RPRoh '(sample|batch)\[".*?\"]' | sort | uniq
batch["boxes"]
batch["image"]
batch["label"]
batch["labels"]
batch["mask"]
batch["prediction"]
batch["prediction_boxes"]
batch["prediction_labels"]
batch["prediction_scores"]
sample["bbox"]
sample["boxes"]
sample["chm"]
sample["hsi"]
sample["image"]
sample["label"]
sample["mask"]
sample["num_patches"]
sample["original_shape"]
sample["padding"]
sample["patch_shape"]
sample["prediction"]
sample["prediction_boxes"]
sample["prediction_label"]

This makes it difficult to create dataset-independent trainers that handle multiple datamodules.

We would also like to be able to type check each of these so we don't have a mix of ints/floats and Tensors (e.g., for label).

TorchGeo relies heavily on Kornia for data augmentation. Kornia provides an AugmentationSequential container for composing and applying transforms. Currently, we create our own wrapper around AugmentationSequential that maps our keys to the keys expected by Kornia, but a better solution would be to standardize on the same keys that Kornia uses.

Implementation

Newer versions of Python introduce the TypedDict type. I propose we use this (and add a dependency on typing-extensions for older versions of Python). The code would look like:

from typing import TypedDict

class Sample(TypedDict, total=False):
    input: Tensor
    label: Tensor
    mask: Tensor
    bbox: Tensor
    bbox_xyxy: Tensor
    bbox_xywh: Tensor
    keypoints: Tensor

Batch = Sample

Each dataset __getitem__ would return an object of type Sample, while collation functions would take in List[Sample] and return Batch. This would replace our current Dict[str, Tensor] type hints and be much more strict.

There are a few important caveats to note here:

1. By specifying total=False, all keys are optional.
2. Extra keys are disallowed. If a dataset requires a custom key, it will need to be added to this list. If this wasn't the case, we would end up with accidental labels and boxes keys again.
3. PyTorch does not yet allow for dtype or shape to be typed. When it does, Batch will have its own separate implementation. For now, the difference is only semantic.

Alternatives

The keys allowed by Kornia are defined in the kornia.constants.DataKey enum. Unfortunately, I couldn't find any documentation or discussion on this enum. Kornia does not standardize "label", so we could alternatively call this "target".

Kornia isn't the only library for data augmentation. There is also Albumentations, which seems to standardize on:

• image
• class_labels
• class_categories
• mask
• bboxes
• bboxes_classes
• keypoints
• keypoints_classes

Torchvision does not use sample dicts, so they don't have a standard of their own. Torchvision uses PIL, while Albumentations uses OpenCV, neither of which support MSI or GPU, so I think we're better off matching the Kornia standard.

The final alternative is to convince Kornia to change their key names to match ours. I think this is unlikely, but @isaaccorley is more optimistic.

Additional information

We will likely need many additional keys for things that don't fall into this list. How strict should we be about this? It would be really nice if PyTorch's default collate function worked. We currently roll our own solely because we need to include things like rasterio.crs.CRS which isn't a type that the default function knows how to handle.

How should we handle predictions? We currently prefix with prediction_.

How should we handle multi-label problems? The dimensions of these would be different. Should we use a different key?

Should we also standardize dtypes? All Kornia transforms require {float16, float32, float64} as input, while all torchmetrics metrics require int? I can't actually find evidence of the latter, but we have several datamodules which assert this. Depending on if labels are for classification or regression, the type can and should change. We could either have separate keys for these, or ignore dtype for these.

[isaaccorley]

I'm okay with standardizing on input, mask, label, bbox, bbox_xyxy, .... I think some datasets have multiple inputs of different shapes and dtypes, e.g. IDTReeS (chm, hsi, rgb, and label). This is the only edge case I can think of at the moment though. There are some other datasets like RSVQA which have text/image pairs which we haven't thought through yet but would be of interest given the current diffusion model hype.

[adamjstewart]

What exactly are the outputs of IDTReeS? I know they have point cloud data so we'll at least need another key for that. Are the rest actually useful? We could always add them as people need them.

[isaaccorley]

The other outputs could be treated like images. HSI is an image of shape (C,H,W). CHM is like a DSM but for canopy height. It's also (C,H,W) but C=1.

[adamjstewart]

Currently, we create our own wrapper around AugmentationSequential that maps our keys to the keys expected by Kornia, but a better solution would be to standardize on the same keys that Kornia uses.

@isaaccorley correctly pointed out that AugmentationSequential doesn't take sample dicts as input and that the user needs to manually pass the input like so:

augs = AugmentationSequential(...)
augs(sample["image"], sample["mask"], data_keys=["input", "mask"])

so there's actually no reason to force ourselves to match our keys to Kornia's.

I think we should:

1. Standardize on "image" for now since that's what we almost exclusively use
2. See if Kornia is willing to support sample dicts as input to AugmentationSequential.forward so that we can do sample = augs(sample): AugmentationSequential: accept sample dict as input kornia/kornia#2119
3. See if Kornia is willing to add "image" as an alias to "input", it's just an enum so it would be a 1-line code change
4. Consider changing "image" to "input" in our code iff 2 is approved but 3 is rejected

Does that sound reasonable? I agree that "input" isn't a good variable name, it's a builtin function in Python and it's too generic.

[MathiasBaumgartinger]

Thanks for your your elaboration! As i think this might be relevant information for other users as well, shall we work together to include this in the correspodning tutorial?

[adamjstewart]

Kornia 0.7.4 is supposed to be released quite soon I think (@johnnv1 @edgarriba?). By the time our tutorial is updated, it will probably be out. @ashnair1 is already working on updating the rest of TorchGeo to be compatible with Kornia and remove the need for any hacks, although some of that work is backwards-incompatible and won't make it into a stable release until TorchGeo 0.7.0. I guess I'll let @ashnair1 decide if it's worth clarifying this in the tutorial or if we should just wait until it's no longer an issue.

[MathiasBaumgartinger]

I see your point. Anyway, I was thinking about a mention of using augs as opposed to transform with the explanation you provided would be helpful. Especially since https://torchgeo.readthedocs.io/en/stable/tutorials/transforms.html does not mention how one would include augmentations in a module at all.

[adamjstewart]

What we really need is an entire tutorial on writing custom data modules.

[edgarriba]

@adamjstewart the reelase target is next week after eccv

[adamjstewart]

We can actually standardize the dtypes at the same time. I propose:

class Sample(TypedDict, total=False):
    input: FloatTensor
    label: LongTensor
    mask: LongTensor
    bbox: LongTensor
    bbox_xyxy: LongTensor
    bbox_xywh: LongTensor
    keypoints: LongTensor

Kornia still requires everything to be float, but torchmetrics and torchvision/timm seem to require all targets to be long.

[adamjstewart]

Or not, see #1149. We may need to see if it's possible for kornia/torchmetrics/torchvision/timm to be less particular about their dtypes.

[adamjstewart]

While we're at it, we should also standardize all transforms/data augmentation:

SampleTransform = Callable[[Sample], Sample]
BatchTransform = Callable[[Batch], Batch]

[adamjstewart]

In #1997 we note that some datasets contain both regression and classification labels, but these can't both be returned since they both use the same "label" key. We should consider using a different key for each type of task such that both classification and regression can be supported for the same dataset.

[adamjstewart]

Extra keys are disallowed. If a dataset requires a custom key, it will need to be added to this list. If this wasn't the case, we would end up with accidental labels and boxes keys again.

We can make subclasses for these: https://stackoverflow.com/a/71814659/5828163

[adamjstewart]

Also possible to disable the specific error raised for these extra keys: python/mypy#4617 (comment). Only problem is that this also requires libraries that depend on us to ignore these mypy warnings too.

[adamjstewart]

Subclasses are annoying because the subclasses must be documented too. I wonder if we can create a custom TypedDict interface with a generic __getitem__(self, key: str) -> Any for all unknown keys.

[MathiasBaumgartinger]

Hey! This is the closest I could find without opening a new discussion.

Is there any way to have some augmentations that are applied on image and mask and others that are only applied on the image?
While transformations (i.e. modifying the spatial dimensions) require to be applied on the mask, modifications on the channel dimension such as ColorJiggle may mess up the training pipeline.

[adamjstewart]

Can you try using only K.ColorJiggle and no other transforms? Also make sure you're using vmin, vmax, and interpolation=None with plt.imshow to make sure it uses the same cmap for all classes.

[MathiasBaumgartinger]

Are you sure this is because of ColorJiggle and not because of the RandomRotate with the background fill=0? Seems like it's decreasing all the mask values -1.

You are right, it does come from the rotation caused by RandomAffine. I was sure this also happened with ColorJiggle only. My bad, Thanks for clarifying.

Have you faced this problem before? Is there any way to circumvent this behavior? Seems buggy to me

[isaaccorley]

I would set the pad_value to the value of your "other" class and see if that works. Not sure what's exactly causing your other classes to be decremented. Would you mind sharing your code?

My guess is that the padding is creating another value outside of your classes and the colormap of the plotting is just being shifted by 1. Printing out the unique values in the mask after you load the batch might help figure this out.

[MathiasBaumgartinger]

Thats a good intuition! The code where all of this happens. Note that FLAIR2 is not yet merged.

def prepare(config: dict, data_dir: str = default_data_dir) -> Tuple[FL, torch.nn.Module]:
    _, _, means, stdvs = FLAIR2.per_band_statistics("train", config["aerial_bands"])

    transforms = AugmentationSequential(
        K.RandomAffine(degrees=(-180, 180), translate=(0.25, 0.25), p=0.5, padding_mode=1, keepdim=True),
        data_keys=None,
        same_on_batch=False,
        keepdim=True,
    ).cuda()

    module = FL(
        num_workers=config["num_workers"],
        batch_size=config["batch_size"],
        patch_size=config["patch_size"],
        val_split_pct=0.2,
        use_toy=True,
        augs=transforms,
        root=data_dir,
        **{"aerial_bands": config["aerial_bands"]},
    )

    task = SemanticSegmentationTask(
        model="unet",
        backbone="resnext101_32x8d",
        ignore_index=255,
        in_channels=len(config["aerial_bands"]),
        num_classes=13,
        lr=config["lr"],
        patience=config["lr_patience"]
    )

    return module, task

def train(
    module: FL,
    task: SemanticSegmentationTask,
    root_dir: str = default_root_dir,
    min_epochs: int = 1,
    max_epochs: int = config["max_epochs"],
) -> Trainer:
    print_setup(assert_setup=True)

    # Callbacks
    checkpoint_callback = ModelCheckpoint(monitor="val_loss", save_top_k=1, mode="min")
    lr_monitor = LearningRateMonitor(logging_interval="step")
    logger = TensorBoardLogger(save_dir=root_dir, name="FLAIR2logs")

    trainer = Trainer(
        accelerator=accelerator,
        num_nodes=1,
        callbacks=[checkpoint_callback, lr_monitor],
        log_every_n_steps=1,
        logger=logger,
        min_epochs=min_epochs,
        max_epochs=max_epochs,
        precision=32,
        enable_progress_bar=True,
        detect_anomaly=True,
    )

    trainer.fit(model=task, datamodule=module)
    
    return trainer


def main():
    module, task = prepare(config)
    trainer = train(module, task)
    trainer.test(model=task, datamodule=module)


if __name__ == "__main__":
    main()

[MathiasBaumgartinger]

Okay, I got it to work; the internal function that is called using K.RandomAffine indeed has a fill mode. https://github.com/kornia/kornia/blob/c4e6bf274f54f9b7475ed5eaf5d7bdd7f7c2d516/kornia/geometry/transform/imgwarp.py#L134-L141

Sadly, neither the padding_mode fill, nor the fill_value are documented. Furthermore, the initialization of the fill_value is not exposed in K.RandomAffine. For anyone facing the same problem, this custom class solves the problem:

from typing import Any, Dict, Optional
import torch
from torch import Tensor
import kornia.augmentation as K
from kornia.geometry.transform import warp_affine


class RandomAffineFill(K.RandomAffine):
    """
    Custom RandomAffine that fixes the output size to be the same as the input size.
    """

    def __init__(self, fill_value: int, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.fill_value = fill_value


    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        _, _, height, width = input.shape
        if not isinstance(transform, Tensor):
            raise TypeError(f"Expected the `transform` be a Tensor. Got {type(transform)}.")

        return warp_affine(
            input,
            transform[:, :2, :],
            (height, width),
            flags["resample"].name.lower(),
            align_corners=flags["align_corners"],
            padding_mode="fill",
            fill_value=torch.tensor([self.fill_value])
        )

EDIT:

Link to Kornia Feature Request: kornia/kornia#3090

[adamjstewart]

If we can't figure out an easy way to make a TypedDict allow additional keys, I'm leaning towards simplifying this to:

Sample = dict[str, Tensor]

This will also guarantee compatibility with PyTorch's builtin default_collate function.

[adamjstewart]

This may be coming in Python 3.14! https://peps.python.org/pep-0728/