Cityscapes#

class torch_uncertainty.datasets.segmentation.Cityscapes(root, split='train', mode='fine', target_type='instance', transform=None, target_transform=None, transforms=None)[source]#

Cityscapes dataset wrapper with train-ID color mapping.

Extends torchvision.datasets.Cityscapes to provide a stable color palette for visualization and convenience helpers to encode and decode semantic segmentation targets using Cityscapes train IDs. A tensor mapping train IDs to RGB colors is constructed on initialization for decoding predicted masks.

Variables:

  • color_palette – List of RGB tuples for each class label in Cityscapes.

  • train_id_to_color – Tensor mapping train IDs to RGB colors for decoding.

Parameters:

  • root (str) – Root directory of the Cityscapes dataset.

  • split (str) – Dataset split to use (e.g., “train”, “val”, or “test”). Defaults to "train".

  • mode (str) – Annotation mode (“fine” or “coarse”). Defaults to "fine".

  • target_type (list[str] | str) – One or more target types to load (“instance”, “semantic”, etc.). Defaults to "instance".

  • transform (Callable[..., Any] | None) – Transformation applied to the input image. Defaults to None.

  • target_transform (Callable[..., Any] | None) – Transformation applied to the target. Defaults to None.

  • transforms (Callable[..., Any] | None) – Combined transformation for image and target. Defaults to None.

decode_target(target)[source]#

Decode a train-ID tensor into an RGB tensor using the palette.

Pixels with value 255 are treated as void and mapped to the last palette entry (black). The returned tensor contains RGB color values for each pixel according to the dataset palette.

Parameters:

target (Tensor) – Integer tensor of train IDs.

Return type:

Tensor

Returns:

A tensor of RGB colors with shape (H, W, 3) mapped from train IDs.

classmethod encode_target(target)[source]#

Encode a Cityscapes target PIL image into a train-ID image.

The input is a color-coded PIL image (train-ID or label colors). The method converts it to a single-channel image where each pixel value is the corresponding train ID.

Parameters:

target (Image) – A PIL image containing the ground-truth target map.

Return type:

Image

Returns:

A PIL image where pixel values correspond to Cityscapes train IDs.

plot_sample(index, ax=None)[source]#

Plot a dataset sample (image and decoded target) for inspection.

Intended behavior: load the sample at index, decode the target to RGB using train_id_to_color, and render the image and overlay/target on the provided axis. If ax is None, a new matplotlib figure and axis should be created.

Parameters:

  • index (int) – Index of the sample to plot.

  • ax (Axes | None) – Optional matplotlib axis to draw on. Defaults to None.

Return type:

tuple[Figure, Union[Axes, ndarray]]

Returns:

A tuple (fig, ax) or the axis object used for plotting, depending on the plotting utility in use.

Raises:

NotImplementedError – This plotting helper is not implemented yet.