---
title: "torchvision classes"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{torchvision classes}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  eval = FALSE
)
```

```{=html}
<style>
.callout-note {  margin: 1em 0;  padding: .6em .9em;  background: #e7f1fd;  border-left: 4px solid #0d6efd;  border-radius: .25rem;}
.callout-note::before {  content: "Note";  display: block;  font-weight: 700;  color: #0a58ca;  margin-bottom: .25rem;}
</style>
```

All outputs of [torchvision datasets](https://torchvision.mlverse.org/reference/#datasets) and [torchvision models](https://torchvision.mlverse.org/reference/#models)
inference results are plain R lists of torch tensors.
But those lists may have very different attributes names and content between object detection and 
semantic segmentation for examples.

In order to ease the correct mapping of transformations and visualization utilities to 
those diverse objects, we have equipped those objects with S3 classes
that will automatically dispatch to the right method, in a similar way as pytorch vision with TVTensors.

The different families of transforms use S3 classes internally to accept objects as 
input, modify output classes for correct pipeline continuation, and dispatch to 
the appropriate drawing helpers.

This article goes through those S3 classes for the user to understand their content and usage.

The code in this vignette is not evaluated when the website is built, since most of 
the code chunks requires a dataset on disk.

```{r setup}
library(torchvision)
library(torch)
```

## What are those S3 classes made for?

A torchvision dataset item is a list of `x`, the tensor image, and `y`, the target that 
is also a list containing `bbox` with a [4, N] tensor in case of object-detection whereas 
it will contain `mask`, a binary tensor in case of a semantic-segmentation item. 

Both the image and its bounding boxes (or mask) must stay consistent : Rotate the image and the boxes
 (respectively mask) have to be rotated as well; crop it and both have to be cropped the same way. 
We do not want to have different function names for each and every transform for each and every object family. So the
S3 class on the item is the placeholder for the computer-vision task family and will 
magically dispatch the transform accordingly.

There is a drawback to this: a dataset item or a dataset target that has missed 
to match one of the defined S3 classes will fail at the next `transform_` call or 
fail to be rendered by drawing helpers like `draw_bounding_boxes(item)`, rather than producing a wrong
result.

## Datasets, items, image and targets

Indexing a dataset returns an *item*: a list with `x`, the *image*, and `y`,
a list of everything the dataset associates to the image. `y` is called the *target*.

![A dataset is a deck of items; each item holds an image in `$x` and its annotations in `$y`.](assets/dataset-item.svg)

A dataset takes up to three transform arguments, and they reach different depths
of the *item*  structure. `transform` is handeling `x` alone, `target_transform` is handeling
`y` alone, and `item_transform` is handeling the item as a whole.

Classification datasets only offer the first two. Detection and
segmentation datasets offer all three, because an image cannot be rotated or
cropped without moving its boxes and masks with it.

```{r}
ds <- coco_detection_dataset(
  train = FALSE,
  download = TRUE,
  transform = \(img) img %>% transform_to_tensor() %>% transform_rgb_to_grayscale(),
  item_transform = \(item) item %>% item_transform_resize(c(480, 640))
)
```

::: {.callout-note}
#### Transform ordering
When multiple transforms are included in a dataset construction, they are applied in the following order:

1. `transform`, applied to image

2. `target_transform` if any, applied to target

3. `item_transform` applied to the results of the two previous ones onto both image and targets.
:::

::: {.callout-note}
 `transform_to_tensor` is doing more work there than it looks. A dataset returns
`x` in whatever form it was read in, which is the `(H, W, 3)` channel-last native R array
out of `base_loader()` . The item transforms and the drawing helpers all want the
`(C, H, W)` tensor that torch models take, so a detection or segmentation
dataset is nearly always built with `transform = transform_to_tensor`. 
:::

The dataset itself gets a S3 class depending on the computer-vision task, either `object_detection_dataset` or
`segmentation_dataset` is inserted ahead of torch generic class `dataset`. That class is what
lets you hand a whole dataset to an item transform:

```{r}
ds <- item_transform_random_horizontal_flip(ds, p = 0.5)
```

which composes the flip into the dataset's `item_transform` rather than applying
it once.

## Classification and caption datasets and their items

The simplest items carry no class at all. `mnist_dataset()`, `cifar10_dataset()`,
`eurosat_dataset()` and the other classification datasets return a bare list
whose `y` is the class index:

```{r}
mnist <- mnist_dataset(root = tempdir(), download = TRUE, transform = transform_to_tensor)
str(mnist[1])
#> List of 2
#>  $ x:Float [1:1, 1:28, 1:28]
#>  $ y: int 6
```

`flickr8k_caption_dataset()` and `flickr30k_caption_dataset()` have the same
shape, with a character vector of captions in `y`. Nothing dispatches on these
items: an image transform goes in through `transform`, and there is no geometry
in the target to keep in sync with it.

## Detection datasets and their items

`coco_detection_dataset()`, `pascal_detection_dataset()` and the
`rf100_*_collection()`s give their items the class `image_with_bounding_box` and
their target the class `object_detection_target`.

![An `image_with_bounding_box` item and the fields of its `object_detection_target`.](assets/detection-item.svg)

| field | contents |
|---|---|
| `$boxes` | `(N, 4)` tensor, one row per object, as `(x_min, y_min, x_max, y_max)` in pixels |
| `$labels` | the class name of each box |
| `$area` | `(N)` float tensor, the area of each annotation |
| `$iscrowd` | `(N)` boolean tensor, `TRUE` where the annotation covers a group of objects rather than one |
| `$image_height`, `$image_width` | the size `x` had when the boxes were read |

Only `$boxes` is guaranteed, since that is the field the class is validated on.
The rest is whatever the annotations happened to carry: the table above is a
COCO target, while an RF100 one holds `$boxes`, `$labels` and an `$image_id` and
nothing else.

`draw_bounding_boxes()` has a method for the item, so it finds `$x`,
`$y$boxes` and `$y$labels` on its own:

```{r}
boxed <- draw_bounding_boxes(ds[1])
tensor_image_browse(boxed)
```

## Rotated box items

Four numbers per box cannot describe a box that has been rotated, so
`item_transform_rotate()` and `item_transform_affine()` return an
`image_with_rotated_box` instead of the item they were given:

```{r}
class(item_transform_rotate(ds[1], angle = 15))
#> [1] "image_with_rotated_box" "list"
```

![An `image_with_rotated_box` item: same target fields, but `$boxes` has a fifth column.](assets/rotated-box-item.svg)

The target keeps its `object_detection_target` class and its fields, but
`$boxes` gains a fifth column holding the accumulated angle in degrees, the
xyxyr format that `box_xyxy_to_xyxyr()` and `target_transform_rotate()` produce.
`draw_bounding_boxes()` has a second method for this class which draws each box
as a polygon through magick rather than as an upright rectangle.

Every item transform has a method for rotated-box items, except
`item_transform_perspective()`.

## Segmentation datasets and their items

`coco_segmentation_dataset()`, `pascal_segmentation_dataset()`,
`cityscapes_dataset()`, `oxfordiiitpet_segmentation_dataset()` and
`rf100_peixos_segmentation_dataset()` return `image_with_segmentation_mask`
items whose target is a `segmentation_target`.

![An `image_with_segmentation_mask` item and its `segmentation_target`.](assets/segmentation-item.svg)

No field is required here. Segmentation annotations differ far more between
datasets than detection ones do, so each dataset passes on what it reads:
polygons in `$segmentation` for COCO, `$polygon`, `$color`, `$semantic` and
`$instance` for Cityscapes depending on the requested `target_type`, one-hot
`$masks` for Oxford-IIIT Pet and Pascal VOC. `$labels` is not uniform either:
the detection datasets name their classes, while
`pascal_segmentation_dataset()` numbers the planes of `$masks` that are not
empty.

What `draw_segmentation_masks()` and the segmentation models want is `$masks`,
an `(N, H, W)` tensor with one plane per annotated class. Two target transforms
build it from raw annotations:

| transform | reads |
|---|---|
| `target_transform_coco_masks()` | `$segmentation`, the COCO polygons |
| `target_transform_trimap_masks()` | `$trimap`, the three-value Oxford-IIIT Pet mask |

```{r}
seg_ds <- coco_segmentation_dataset(
  train = FALSE,
  download = TRUE,
  transform = transform_to_tensor,
  target_transform = target_transform_coco_masks
)
masked <- draw_segmentation_masks(seg_ds[1])
```

`coco_segmentation_dataset()` only promotes its items to
`image_with_segmentation_mask` once `$masks` is there, so without that target
transform they stay `image_with_bounding_box`. `cityscapes_dataset()` has no
target transform of its own, and its targets have to be turned into `$masks` by
hand before they can be drawn.

## Slicing an image: `sahi_split`

SAHI, Slicing Aided Hyper Inference, cuts a large image into overlapping tiles,
runs the model on each and puts the results back together. The tiling itself is
an object: `prepare_sahi_split()` measures an image and returns a `sahi_split`.

![A `sahi_split`, drawn as one card rather than a deck, because it is computed once and reused across images.](assets/sahi-split.svg)

```{r}
split <- prepare_sahi_split(c(1000, 1200), size = c(512, 512))
length(split$crop_windows)
#> [1] 9
```

It holds `$crop_windows`, one `top`/`left`/`height`/`width` window per tile, and
the `$size`, `$overlap_size_ratio`, `$image_height` and `$image_width` it was
computed from. Keeping it as its own object is what lets the same tiling be
applied to an image and to its target, so the two stay aligned:

```{r}
tiled <- ds[1]
split <- prepare_sahi_split(tiled$x, size = c(256, 256))

crops <- transform_sahi_crop(tiled$x, split)
targets <- target_transform_sahi_crop(tiled$y, split)
```

`transform_sahi_crop()` returns the tiles stacked as an `(N, C, H, W)` batch.
`target_transform_sahi_crop()` returns one target per tile, each box clipped to
the tile and shifted into its coordinates, and `$image_height` and
`$image_width` set to the size of the tile rather than of the original image. A
box is kept only if what survives the clip is at least `min_area_ratio` of its
original area, a tenth by default, so a box grazing the edge of a tile is
dropped from that tile rather than kept as a sliver.

Unlike the item transforms, this pair turns one item into many, so there is no
`item_transform_sahi_crop()` doing both halves at once.

## Which function takes which class

| family | dispatches on | acts on |
|---|---|---|
| `transform_*()` | `torch_tensor` | the image alone |
| `prepare_sahi_split()` | `numeric`, `torch_tensor`, `magick-image`, `dataset` | measures an image, returns a `sahi_split` |
| `target_transform_resize()`, `target_transform_rotate()`, `target_transform_affine()`, `target_transform_sahi_crop()` | `object_detection_target` | the target alone |
| `target_transform_coco_masks()`, `target_transform_trimap_masks()` | `segmentation_target` | the target alone |
| `item_transform_*()` | `image_with_bounding_box`, `image_with_rotated_box`, `image_with_segmentation_mask`, `dataset` | image and target together |
| `draw_bounding_boxes()` | `torch_tensor`, `image_with_bounding_box`, `image_with_rotated_box` | draws boxes on the image |
| `draw_segmentation_masks()` | `torch_tensor`, `image_with_segmentation_mask` | overlays masks on the image |

The two groups of target transforms are disjoint, which is why they are two rows
and not one. The first group moves boxes and is written for detection targets
only; the second builds `$masks` out of the annotations a segmentation dataset
carries. Crossing them is an error rather than a silent no-op, so
`target_transform_rotate()` on a `segmentation_target`, or
`target_transform_coco_masks()` on an `object_detection_target`, stops with
`not implemented for` and the class it was handed.

Both groups also accept bare lists with the fields their transforms require.
Transforming a bare list preserves its list class.

The image transforms are written for `torch_tensor`. Five of them also take a
magick image, `transform_crop()`, `transform_hflip()`, `transform_resize()`,
`transform_sahi_crop()` and `transform_to_tensor()`, and `transform_to_tensor()`
alone takes an `array`, a `matrix` or a `list`, which is how an image coming off
a loader gets into a tensor in the first place.

The item transforms are the most uniform family: each of them takes all three
item classes and a `dataset`, apart from `item_transform_perspective()`, which
has no rotated-box method. Among the target transforms only
`target_transform_rotate()` takes a dataset as well. The reference index lists
each of these families in full.

`draw_keypoints()` is not generic: it takes an image tensor and an `(N, K, 2)`
keypoint tensor.

Calling any of these on something they have no method for stops with an error
naming the class it was given, which is usually a sign that a target lost its
class somewhere along a pipeline.

## Building an item by hand

Target transforms accept a bare list of annotations. When building a complete
item for item transforms or drawing functions, give the item its task class;
the example also gives the target its corresponding class:

```{r}
img <- torch_randint(0, 255, size = c(3, 120, 160))$to(dtype = torch_uint8())

y <- list(
  boxes = torch_tensor(rbind(c(10, 20, 60, 90), c(70, 30, 150, 110)), dtype = torch_float()),
  labels = c("cat", "dog"),
  image_height = 120,
  image_width = 160
)
class(y) <- c("object_detection_target", "list")

item <- list(x = img, y = y)
class(item) <- c("image_with_bounding_box", "list")

str(item)
#> List of 2
#>  $ x:Byte [1:3, 1:120, 1:160]
#>  $ y:List of 4
#>   ..$ boxes       :Float [1:2, 1:4]
#>   ..$ labels      : chr [1:2] "cat" "dog"
#>   ..$ image_height: num 120
#>   ..$ image_width : num 160
#>   ..- attr(*, "class")= chr [1:2] "object_detection_target" "list"
#>  - attr(*, "class")= chr [1:2] "image_with_bounding_box" "list"
```

A segmentation target is built the same way, with `segmentation_target` and
`image_with_segmentation_mask`.

## What models return

Model outputs are plain lists and carry no class of their own. The object
detection models return `$detections`, one element per image of the batch, each
with `$boxes`, `$labels` and `$scores`; the semantic segmentation models return
`$out`, the per-class logits. The
[Visualization utilities](visualization-utilities.html) article covers
turning either of them into something the drawing helpers can render.

## Comparison with torchvision TVTensors (advanced)

If you have used the v2 transforms of the Python torchvision, these classes
cover the same ground as
[TVTensors](https://docs.pytorch.org/vision/main/auto_examples/transforms/plot_transforms_getting_started.html#videos-boxes-masks-keypoints),
but they are attached one level further out.

A TVTensor is a subclass of `torch.Tensor`. `Image`, `Video`, `BoundingBoxes`,
`Mask` and `KeyPoints` are tensors that know what they are, and a
`BoundingBoxes` carries its `format` and its `canvas_size` as attributes of the
tensor itself. A v2 transform walks an arbitrary nested sample, a dict, a list
or a tuple, dispatches on the type of each element it finds, hands back the same
structure, and passes anything it does not recognise through untouched.

Here the class goes on the list instead. `$x` and `$y$boxes` are ordinary
`torch_tensor`s, and what carries the meaning is `image_with_bounding_box` on
the item and `object_detection_target` on `y`. Four differences follow from
that:

- An item is always `list(x, y)` rather than any nesting you like, and the three
  transform families say which half of it they act on instead of one transform
  traversing the whole sample.
- Box format is not an attribute but the width of `$boxes`: four columns are
  xyxy, five are xyxyr. What `canvas_size` holds lives in `$image_height` and
  `$image_width` on the target.
- A change of geometry changes the class of the item, so
  `item_transform_rotate()` hands back an `image_with_rotated_box`. On the
  Python side the object stays a `BoundingBoxes` and the distinction lives in
  its `format`.
- There is no heuristic for unwrapped input. A v2 transform whose sample holds
  no `Image` will treat the first plain tensor it finds as the image; here a
  list without its class is an error, not a guess.

Two of the five TVTensors have no counterpart here. There is no video support at
all, and keypoints are not a class: `draw_keypoints()` takes an image tensor and
a keypoint tensor and nothing dispatches on either.