Datasets: always return sample = dict[str, Tensor]

[adamjstewart]

This PR introduces a new Sample type alias defined as:

Sample = dict[str, Tensor]

All dataset __getitem__ methods return a Sample, and all transforms operate on a Sample.

This PR supersedes #2249 with a simpler approach (can be expanded later).

Pros

• Compatibility with PyTorch's default_collate function
• Compatibility with Lightning's transfer_batch_to_device
• Uniform type hints across all TorchGeo subpackages
• Avoids most uses of dynamic typing (typing.Any)
• Allows arbitrary key names (unlike Add Sample and Batch TypedDicts #2249)

Cons

• Datasets can no longer return non-Tensor types
• Does not enforce standard key names (unlike Add Sample and Batch TypedDicts #2249)
• Backwards-incompatible changes (removing CRS/bounds)

Closes #2249
See #985 for discussion

[adamjstewart]

@ashnair1 do object detection bounding boxes have to be list[Tensor] or can they be Tensor? I wonder if we can use some kind of packed/padded sequence logic.

[ashnair1]

@ashnair1 do object detection bounding boxes have to be list[Tensor] or can they be Tensor? I wonder if we can use some kind of packed/padded sequence logic.

Object detection bboxes are usually list[Tensor]. This is because each image can have a varying number of objects e.g.

img1 -> [5, 4], 
img2 -> [15, 4], 
img3 -> [1, 4] 

Not sure about packing/padding. I think given a batch, we could theoretically iterate over the box tensors, find the max and pad the rest accordingly. So above example would be

img1 -> [15, 4], # Was 5, padded to 15 
img2 -> [15, 4],  # 15 max
img3 -> [15, 4], # Was 1, padded to 15

But I don't think it's optimal.

[adamjstewart]

In order to create mini-batches, do we need to pad anyway? PyTorch has a PackedSequence class but it's still not a subclass of Tensor. We might be back to #2249 if we need to support dicts with multiple value types. But that doesn't support unknown keys...