peal.data.custom_datasets.SquareDataset

class peal.data.custom_datasets.SquareDataset

Bases: Image2MixedDataset

Synthetic 64x64 squares whose colour is confounded by the background.

An image shows a bordered square with red interior intensity ColorA (the true feature) on a grey background of intensity ColorB (the confounder), at one of nine positions. SquareDatasetGenerator writes the dataset on first use and only supports config.confounder_probability == 0.5; confounded variants are derived from that pool elsewhere. Since both generative factors are known, this class provides an exact 2D latent readout for the decision boundary and counterfactual plots, either from the hint mask or from an oracle model.

Parameters:

• config (DataConfig) – Uses dataset_path, confounder_probability and confounding_factors.

• **kwargs – Forwarded to Image2MixedDataset.

Variables:

oracle (torch.nn.Module, optional) – $PEAL_RUNS/square/colora_confounding_colorb/ classifier_all_attributes_resnet18/model.cpl in eval mode, if present. Used by sample_to_2d_latent when no mask is given.

Raises:

NotImplementedError – If the dataset has to be generated but confounder_probability is not 0.5.

Methods

check_background	Mean intensity outside the hint mask, i.e. the background level.
check_foreground	Mean red intensity inside the hint mask, i.e. the square's interior.
generate_contrastive_collage	Contrastive collage annotated with the oracle and hint latents.
global_counterfactual_visualization	Render the counterfactual arrow plot twice: from the oracle and the hints.
sample_to_2d_latent	Map a sample to its two generative factors.
visualize_decision_boundary	Draw both the generic oracle plot and the precise hint-based plot.
visualize_decision_boundary_precise	Render the decision boundary over the true generative factors.

Attributes

__init__(config, **kwargs)

Generate the dataset if missing and load the latent oracle if present.

Parameters:

config (DataConfig)

visualize_decision_boundary(*args, **kwargs)

Draw both the generic oracle plot and the precise hint-based plot.

The base implementation is run with hints disabled (so the predictor sees plain images), the previous hint setting is restored, and visualize_decision_boundary_precise is called with the same arguments.

visualize_decision_boundary_precise(predictor, batch_size, device, path, temperature=1.0, train_dataloader=None, val_dataloaders=[], val_weights=[], test_dataloader=None, **kwargs)

Render the decision boundary over the true generative factors.

Unlike the base implementation, which fits a model in an estimated latent space, this evaluates predictor on images synthesised directly from the grid: for each of the nine square positions a 100x100 grid of (ColorA, ColorB) pairs is rendered with latent_to_square_image and the softmax probability of class 0 is averaged over the positions. The grid is cached next to the plot as a .npy file and reused on later calls. Train and validation samples are then overlaid at their measured foreground and background intensity, which needs hints and therefore temporarily enables them on the given loaders.

Parameters:

• predictor (torch.nn.Module) – Classifier, evaluated on device in batches of batch_size.

• batch_size (int) – Batch size of the grid evaluation.

• device (torch.device or str) – Device the predictor runs on.

• path (str) – Path of the non-precise plot; .png is replaced by _precise.png and the cached grid by _precise.npy.

• temperature (float, optional) – Softmax temperature applied to the logits.

• train_dataloader (DataLoader or DataloaderMixer, optional) – Up to 8 batches are plotted as triangles.

• val_dataloaders (list or WeightedDataloaderList, optional) – Validation loaders, plotted as squares; each gets a budget of 8 // val_weights[i] batches.

• val_weights (list of float, optional) – Weights of val_dataloaders; taken from the list itself when it is a WeightedDataloaderList.

• test_dataloader (DataLoader, optional) – Unused; kept for interface compatibility.

check_foreground(x, hint)

Mean red intensity inside the hint mask, i.e. the square’s interior.

Parameters:

• x (torch.Tensor) – Image [..., C, H, W] in the dataset’s own value range; only channel 0 is read.

• hint (torch.Tensor) – Binary mask of the inner square with the same shape.

Returns:

Scalar mean of the masked red channel.

Return type:

torch.Tensor

Notes

In contrast to check_background the image is not projected to the pytorch default range first.

check_background(x, hint)

Mean intensity outside the hint mask, i.e. the background level.

x is mapped to the [0, 1] pytorch default range first and averaged over all channels where the mask is zero.

Parameters:

• x (torch.Tensor) – Image [..., C, H, W] in the dataset’s own value range.

• hint (torch.Tensor) – Binary mask of the inner square.

Returns:

Scalar mean intensity of the unmasked area.

Return type:

torch.Tensor

sample_to_2d_latent(sample, mask=None)

Map a sample to its two generative factors.

With a hint mask the factors are measured directly as [check_foreground, check_background]. Without a mask the oracle classifier is applied instead and, if config.confounding_factors is set, its outputs are reduced to the attributes named there.

Parameters:

• sample (torch.Tensor) – Image [C, H, W] or batch [N, C, H, W]; a single image is unsqueezed for the oracle and squeezed again afterwards.

• mask (torch.Tensor, optional) – Hint mask of the inner square.

Returns:

The two latent values.

Return type:

torch.Tensor

global_counterfactual_visualization(filename, x_list, counterfactuals, y_target_start_confidence, y_target_end_confidence, y_list, hint_list, attempts=1)

Render the counterfactual arrow plot twice: from the oracle and the hints.

The base implementation is first called with hints disabled and without masks (latents from the oracle), then again with hint_list writing to *_precise.png - matching the two grids that visualize_decision_boundary saves.

Parameters:

• filename (str) – Output path of the oracle version; the precise one replaces .png by _precise.png.

• x_list (list of torch.Tensor) – Factual and counterfactual images.

• counterfactuals (list of torch.Tensor) – Factual and counterfactual images.

• y_target_start_confidence (list of float) – Target-class confidence before and after the edit.

• y_target_end_confidence (list of float) – Target-class confidence before and after the edit.

• y_list (list) – Ground truth labels.

• hint_list (list of torch.Tensor) – Masks used by the precise version.

• attempts (int, optional) – Counterfactual attempts per sample (one subplot each).

generate_contrastive_collage(x_list, x_counterfactual_list, y_target_list, y_source_list, y_list, y_target_start_confidence_list, y_target_end_confidence_list, base_path, idx_to_info=None, **kwargs)

Contrastive collage annotated with the oracle and hint latents.

When no idx_to_info callback is passed, one is built that reports per confounding factor the factual and counterfactual value as measured by the oracle and by the hint mask, e.g. ColorA: 0.2 (Oracle) / 0.21 (Hint) -> 0.8 (Oracle) / 0.79 (Hint). Everything else is delegated to ImageDataset.generate_contrastive_collage.

Parameters:

• x_list (list of torch.Tensor) – Factual and counterfactual images.

• x_counterfactual_list (list of torch.Tensor) – Factual and counterfactual images.

• y_target_list (list) – Target class, predicted source class and ground truth per pair.

• y_source_list (list) – Target class, predicted source class and ground truth per pair.

• y_list (list) – Target class, predicted source class and ground truth per pair.

• y_target_start_confidence_list (list) – Target-class confidence before and after the edit.

• y_target_end_confidence_list (list) – Target-class confidence before and after the edit.

• base_path (str) – Output directory of the collages.

• idx_to_info (callable, optional) – (x, x_counterfactual, hint) -> str used for the title suffix.

• **kwargs – Passed through to the base implementation.

Returns:

Whatever the base implementation returns.

Return type:

tuple