equiadapt.images package

Subpackages

Submodules

equiadapt.images.utils module

equiadapt.images.utils.flip_boxes(boxes: Tensor, width: int) → Tensor[source]

Flips bounding boxes horizontally.

Parameters:

  • boxes (torch.Tensor) – The bounding boxes to flip.

  • width (int) – The width of the image.

Returns:

The flipped bounding boxes.

Return type:

torch.Tensor

equiadapt.images.utils.flip_masks(masks: Tensor) → Tensor[source]

Flips masks horizontally.

Parameters:

masks (torch.Tensor) – The masks to flip.

Returns:

The flipped masks.

Return type:

torch.Tensor

equiadapt.images.utils.get_action_on_image_features(feature_map: Tensor, group_info_dict: dict, group_element_dict: dict, induced_rep_type: str = 'regular') → Tensor[source]

Applies a group action to the feature map.

Parameters:

  • feature_map (torch.Tensor) – The input feature map.

  • group_info_dict (dict) – A dictionary containing information about the group.

  • group_element_dict (dict) – A dictionary containing the group elements.

  • induced_rep_type (str, optional) – The type of induced representation. Defaults to “regular”.

Returns:

The feature map after the group action has been applied.

Return type:

torch.Tensor

equiadapt.images.utils.roll_by_gather(feature_map: Tensor, shifts: Tensor) → Tensor[source]

Shifts the feature map along the group dimension by the specified shifts.

Parameters:

  • feature_map (torch.Tensor) – The input feature map. It should have the shape (batch, channel, group, x_dim, y_dim).

  • shifts (torch.Tensor) – The shifts for each feature map in the batch.

Returns:

The shifted feature map.

Return type:

torch.Tensor

equiadapt.images.utils.rotate_boxes(boxes: Tensor, angle: Tensor, width: int) → Tensor[source]

Rotates bounding boxes by a specified angle.

Parameters:

  • boxes (torch.Tensor) – The bounding boxes to rotate.

  • angle (torch.Tensor) – The angle to rotate the bounding boxes by.

  • width (int) – The width of the image.

Returns:

The rotated bounding boxes.

Return type:

torch.Tensor

equiadapt.images.utils.rotate_masks(masks: Tensor, angle: Tensor) → Tensor[source]

Rotates masks by a specified angle.

Parameters:

  • masks (torch.Tensor) – The masks to rotate.

  • angle (torch.Tensor) – The angle to rotate the masks by.

Returns:

The rotated masks.

Return type:

torch.Tensor

equiadapt.images.utils.rotate_points(origin: List[float], point: Tensor, angle: Tensor) → Tuple[Tensor, Tensor][source]

Rotates a point around an origin by a specified angle.

Parameters:

  • origin (List[float]) – The origin to rotate the point around.

  • point (torch.Tensor) – The point to rotate.

  • angle (torch.Tensor) – The angle to rotate the point by.

Returns:

The rotated point.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Module contents

class equiadapt.images.ContinuousGroupImageCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig, in_shape: tuple)[source]

Bases: ContinuousGroupCanonicalization

This class represents a continuous group image canonicalization model.

The model is designed to be equivariant under a continuous group of transformations, which can include rotations and reflections. Other specific continuous group image canonicalization classes can be derived from this class.

__init__()[source]

Initializes the ContinuousGroupImageCanonicalization instance.

get_rotation_matrix_from_vector()

This method takes the input vector and returns the rotation matrix.

get_groupelement()[source]

This method maps the input image to the group element.

transformations_before_canonicalization_network_forward()[source]

Applies transformations to the input image before forwarding it through the canonicalization network.

get_group_from_out_vectors()[source]

This method takes the output of the canonicalization network and returns the group element.

canonicalize()[source]

This method takes an image as input and returns the canonicalized image.

invert_canonicalization()[source]

Inverts the canonicalization process on the output of the canonicalized image.

canonicalize(x: Tensor, targets: List | None = None, **kwargs: Any) → Tensor | Tuple[Tensor, List][source]

This method takes an image as input and returns the canonicalized image

Parameters:

  • x (torch.Tensor) – The input image.

  • targets (Optional[List]) – The targets, if any.

Returns:

canonicalized image

Return type:

torch.Tensor

get_group_from_out_vectors(out_vectors: Tensor) → Tuple[dict, Tensor][source]

This method takes the output of the canonicalization network and returns the group element

Parameters:

out_vectors (torch.Tensor) – output of the canonicalization network

Returns:

group element torch.Tensor: group element representation

Return type:

dict

get_groupelement(x: Tensor) → dict[source]

This method takes the input image and maps it to the group element

Parameters:

x (torch.Tensor) – input image

Returns:

group element

Return type:

dict

invert_canonicalization(x_canonicalized_out: Tensor, **kwargs: Any) → Tensor[source]

Inverts the canonicalization process on the output of the canonicalized image.

Parameters:

x_canonicalized_out (torch.Tensor) – The output of the canonicalized image.

Returns:

The output corresponding to the original image.

Return type:

torch.Tensor

transformations_before_canonicalization_network_forward(x: Tensor) → Tensor[source]

Applies transformations to the input image before forwarding it through the canonicalization network.

Parameters:

x (torch.Tensor) – The input image.

Returns:

The transformed image.

Return type:

torch.Tensor

class equiadapt.images.ConvNetwork(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)[source]

Bases: Module

This class represents a convolutional neural network.

The network consists of a sequence of convolutional layers, each followed by batch normalization and a GELU activation function. The number of output channels of the convolutional layers increases after every third layer. The network ends with a fully connected layer.

__init__()[source]

Initializes the ConvNetwork instance.

forward()[source]

Performs a forward pass through the network.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, out_vector_size).

Return type:

torch.Tensor

class equiadapt.images.CustomEquivariantNetwork(in_shape: Tuple[int, int, int, int], out_channels: int, kernel_size: int, group_type: str = 'rotation', num_rotations: int = 4, num_layers: int = 1, device: str = 'cpu')[source]

Bases: Module

This class represents a custom equivariant network.

The network is equivariant to a specified group, which can be either the rotation group or the roto-reflection group. The network consists of a sequence of equivariant convolutional layers, each followed by a ReLU activation function.

__init__()[source]

Initializes the CustomEquivariantNetwork instance.

forward()[source]

Performs a forward pass through the network.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, group_size).

Return type:

torch.Tensor

class equiadapt.images.DiscreteGroupImageCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig, in_shape: tuple)[source]

Bases: DiscreteGroupCanonicalization

This class represents a discrete group image canonicalization model.

The model is designed to be equivariant under a discrete group of transformations, which can include rotations and reflections. Other discrete group canonicalizers can be derived from this class.

__init__()[source]

Initializes the DiscreteGroupImageCanonicalization instance.

groupactivations_to_groupelement()[source]

Takes the activations for each group element as input and returns the group element.

get_groupelement()[source]

Maps the input image to a group element.

transformations_before_canonicalization_network_forward()[source]

Applies transformations to the input images before passing it through the canonicalization network.

canonicalize()[source]

Canonicalizes the input images.

invert_canonicalization()[source]

Inverts the canonicalization of the output of the canonicalized image.

canonicalize(x: Tensor, targets: List | None = None, **kwargs: Any) → Tensor | Tuple[Tensor, List][source]

Canonicalizes the input images.

Parameters:

  • x (torch.Tensor) – The input images.

  • targets (Optional[List], optional) – The targets for instance segmentation. Defaults to None.

  • **kwargs (Any) – Additional keyword arguments.

Returns:

The canonicalized image, and optionally the targets.

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, List]]

get_group_activations(x: Tensor) → Tensor[source]

Gets the group activations for the input images.

Parameters:

x (torch.Tensor) – The input images.

Returns:

The group activations.

Return type:

torch.Tensor

get_groupelement(x: Tensor) → Dict[str, Tensor][source]

Maps the input image to a group element.

Parameters:

x (torch.Tensor) – The input images.

Returns:

The corresponding group elements.

Return type:

dict[str, torch.Tensor]

groupactivations_to_groupelement(group_activations: Tensor) → dict[source]

This method takes the activations for each group element as input and returns the group element

Parameters:

group_activations (torch.Tensor) – activations for each group element.

Returns:

group element.

Return type:

dict

invert_canonicalization(x_canonicalized_out: Tensor, **kwargs: Any) → Tensor[source]

Inverts the canonicalization of the output of the canonicalized image.

Parameters:

  • x_canonicalized_out (torch.Tensor) – The output of the canonicalized image.

  • **kwargs (Any) – Additional keyword arguments.

Returns:

The output corresponding to the original image.

Return type:

torch.Tensor

transformations_before_canonicalization_network_forward(x: Tensor) → Tensor[source]

Applies transformations to the input images before passing it through the canonicalization network.

Parameters:

x (torch.Tensor) – The input image.

Returns:

The pre-canonicalized image.

Return type:

torch.Tensor

class equiadapt.images.ESCNNEquivariantNetwork(in_shape: tuple, out_channels: int, kernel_size: int, group_type: str = 'rotation', num_rotations: int = 4, num_layers: int = 1)[source]

Bases: Module

This class represents an Equivariant Convolutional Neural Network (Equivariant CNN).

The network is equivariant to a group of transformations, which can be either rotations or roto-reflections. The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization, a ReLU activation function, and dropout. The number of output channels of the convolutional layers is the same for all layers.

__init__()[source]

Initializes the ESCNNEquivariantNetwork instance.

forward()[source]

Performs a forward pass through the network.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, num_group_elements).

Return type:

torch.Tensor

class equiadapt.images.ESCNNSteerableNetwork(in_shape: tuple, out_channels: int, kernel_size: int = 9, group_type: str = 'rotation', num_layers: int = 1)[source]

Bases: Module

This class represents a Steerable Equivariant Convolutional Neural Network (Equivariant CNN).

The network is equivariant under all planar rotations. The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a FourierELU activation function. The number of output channels of the convolutional layers is the same for all layers.

__init__()[source]

Initializes the ESCNNSteerableNetwork instance.

forward()[source]

Performs a forward pass through the network.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 2, 2).

Return type:

torch.Tensor

class equiadapt.images.ESCNNWRNEquivariantNetwork(in_shape: tuple, out_channels: int = 64, kernel_size: int = 9, group_type: str = 'rotation', num_layers: int = 12, num_rotations: int = 4)[source]

Bases: Module

This class represents a Wide Residual Network (WRN) that is equivariant under rotations or roto-reflections.

The network consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()[source]

Initializes the ESCNNWRNEquivariantNetwork instance.

forward()[source]

Performs a forward pass through the network.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, group_size).

Return type:

torch.Tensor

class equiadapt.images.ESCNNWideBasic(in_type: FieldType, middle_type: FieldType, out_type: FieldType, kernel_size: int = 3)[source]

Bases: EquivariantModule

This class represents a wide basic layer for an Equivariant Convolutional Neural Network (Equivariant CNN).

The layer consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()[source]

Initializes the ESCNNWideBasic instance.

forward()[source]

Performs a forward pass through the layer.

evaluate_output_shape(input_shape: Tuple[int]) → Tuple[int][source]

Compute the shape the output tensor which would be generated by this module when a tensor with shape input_shape is provided as input.

Parameters:

input_shape (tuple) – shape of the input tensor

Returns:

shape of the output tensor

forward(x: GeometricTensor) → GeometricTensor[source]

Performs a forward pass through the layer.

Parameters:

x (e2cnn.nn.GeometricTensor) – The input data.

Returns:

The output of the layer. The input is added to the output (residual connection).

Return type:

e2cnn.nn.GeometricTensor

class equiadapt.images.ESCNNWideBottleneck(in_type: FieldType, middle_type: FieldType, out_type: FieldType, kernel_size: int = 3)[source]

Bases: EquivariantModule

This class represents a wide bottleneck layer for an Equivariant Convolutional Neural Network (Equivariant CNN).

The layer consists of a sequence of equivariant convolutional layers, each followed by batch normalization and a ReLU activation function. The number of output channels of the convolutional layers is the same for all layers. The input is added to the output of the layer (residual connection).

__init__()[source]

Initializes the ESCNNWideBottleneck instance.

forward()[source]

Performs a forward pass through the layer.

evaluate_output_shape(input_shape: Tuple[int]) → Tuple[int][source]

Compute the shape the output tensor which would be generated by this module when a tensor with shape input_shape is provided as input.

Parameters:

input_shape (tuple) – shape of the input tensor

Returns:

shape of the output tensor

forward(x: GeometricTensor) → GeometricTensor[source]

Performs a forward pass through the layer.

Parameters:

x (e2cnn.nn.GeometricTensor) – The input data.

Returns:

The output of the layer. The input is added to the output (residual connection).

Return type:

e2cnn.nn.GeometricTensor

class equiadapt.images.GroupEquivariantImageCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig, in_shape: tuple)[source]

Bases: DiscreteGroupImageCanonicalization

This class represents a discrete group equivariant image canonicalization model.

The model is designed to be equivariant under a discrete group of transformations, which can include rotations and reflections.

__init__()[source]

Initializes the GroupEquivariantImageCanonicalization instance.

get_group_activations()[source]

Gets the group activations for the input images.

get_group_activations(x: Tensor) → Tensor[source]

Gets the group activations for the input image.

This method takes an image as input, applies transformations before forwarding it through the canonicalization network, and then returns the group activations.

Parameters:

x (torch.Tensor) – The input image.

Returns:

The group activations.

Return type:

torch.Tensor

class equiadapt.images.OptimizedGroupEquivariantImageCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig, in_shape: tuple)[source]

Bases: DiscreteGroupImageCanonicalization

This class represents an optimized (discrete) group equivariant image canonicalization model.

The model is designed to be equivariant under a discrete group of transformations, which can include rotations and reflections.

__init__()[source]

Initializes the OptimizedGroupEquivariantImageCanonicalization instance.

rotate_and_maybe_reflect()[source]

Rotate and maybe reflect the input images.

group_augment()[source]

Augment the input images by applying group transformations (rotations and reflections).

get_group_activations()[source]

Gets the group activations for the input images.

get_optimization_specific_loss()[source]

Gets the loss specific to the optimization process.

get_group_activations(x: Tensor) → Tensor[source]

Gets the group activations for the input image.

Parameters:

x (torch.Tensor) – The input image.

Returns:

The group activations.

Return type:

torch.Tensor

get_optimization_specific_loss() → Tensor[source]

Gets the loss specific to the optimization process.

Returns:

The loss.

Return type:

torch.Tensor

group_augment(x: Tensor) → Tensor[source]

Augment the input images by applying group transformations (rotations and reflections).

Parameters:

x (torch.Tensor) – The input image.

Returns:

The augmented image.

Return type:

torch.Tensor

rotate_and_maybe_reflect(x: Tensor, degrees: Tensor, reflect: bool = False) → List[Tensor][source]

Rotate and maybe reflect the input images.

Parameters:

  • x (torch.Tensor) – The input image.

  • degrees (torch.Tensor) – The degrees of rotation.

  • reflect (bool, optional) – Whether to reflect the image. Defaults to False.

Returns:

The list of rotated and maybe reflected images.

Return type:

List[torch.Tensor]

class equiadapt.images.OptimizedSteerableImageCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig, in_shape: tuple)[source]

Bases: ContinuousGroupImageCanonicalization

This class represents an optimized steerable image canonicalization model.

The model is designed to be equivariant under a continuous group of transformations, which can include rotations and reflections.

__init__()[source]

Initializes the OptimizedSteerableImageCanonicalization instance.

get_rotation_matrix_from_vector()[source]

This method takes the input vector and returns the rotation matrix.

group_augment()[source]

This method applies random rotations and reflections to the input images.

get_groupelement()[source]

This method maps the input image to the group element.

get_optimization_specific_loss()[source]

This method returns the optimization specific loss.

get_groupelement(x: Tensor) → dict[source]

Maps the input image to the group element.

Parameters:

x (torch.Tensor) – The input image.

Returns:

The group element.

Return type:

dict

get_optimization_specific_loss() → Tensor[source]

This method returns the optimization specific loss

Returns:

optimization specific loss

Return type:

torch.Tensor

get_rotation_matrix_from_vector(vectors: Tensor) → Tensor[source]

This method takes the input vector and returns the rotation matrix

Parameters:

vectors (torch.Tensor) – input vector

Returns:

rotation matrices

Return type:

torch.Tensor

group_augment(x: Tensor) → Tuple[Tensor, Tensor][source]

Augmentation of the input images by applying random rotations and, if applicable, reflections, with corresponding transformation matrices.

Parameters:

x (torch.Tensor) – Input images of shape (batch_size, in_channels, height, width).

Returns:

Augmented images. torch.Tensor: Corresponding transformation matrices.

Return type:

torch.Tensor

class equiadapt.images.ResNet18Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)[source]

Bases: Module

This class represents a neural network based on the ResNet-18 architecture.

The network uses a pre-trained ResNet-18 model without its weights. The final fully connected layer of the ResNet-18 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

class equiadapt.images.RotationEquivariantConv(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')[source]

Bases: Module

This class represents a rotation equivariant convolutional layer.

The layer is equivariant to a group of rotations. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()[source]

Initializes the RotationEquivariantConv instance.

get_rotated_permuted_weights()[source]

Returns the weights of the layer after rotation and permutation.

forward()[source]

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, num_rotations, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_rotations, height, width).

Return type:

torch.Tensor

get_rotated_permuted_weights(weights: Tensor, num_rotations: int = 4) → Tensor[source]

Returns the weights of the layer after rotation and permutation.

Parameters:

  • weights (torch.Tensor) – The weights of the layer.

  • num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation and permutation.

Return type:

torch.Tensor

class equiadapt.images.RotationEquivariantConvLift(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')[source]

Bases: Module

This class represents a rotation equivariant convolutional layer with lifting.

The layer is equivariant to a group of rotations. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()[source]

Initializes the RotationEquivariantConvLift instance.

get_rotated_weights()[source]

Returns the weights of the layer after rotation.

forward()[source]

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_rotations, height, width).

Return type:

torch.Tensor

get_rotated_weights(weights: Tensor, num_rotations: int = 4) → Tensor[source]

Returns the weights of the layer after rotation.

Parameters:

  • weights (torch.Tensor) – The weights of the layer.

  • num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation.

Return type:

torch.Tensor

class equiadapt.images.RotoReflectionEquivariantConv(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')[source]

Bases: Module

This class represents a roto-reflection equivariant convolutional layer.

The layer is equivariant to a group of rotations and reflections. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()[source]

Initializes the RotoReflectionEquivariantConv instance.

get_rotoreflected_permuted_weights()[source]

Returns the weights of the layer after rotation, reflection, and permutation.

forward()[source]

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, num_group_elements, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_group_elements, height, width).

Return type:

torch.Tensor

get_rotoreflected_permuted_weights(weights: Tensor, num_rotations: int = 4) → Tensor[source]

Returns the weights of the layer after rotation, reflection, and permutation.

Parameters:

  • weights (torch.Tensor) – The weights of the layer.

  • num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation, reflection, and permutation.

Return type:

torch.Tensor

class equiadapt.images.RotoReflectionEquivariantConvLift(in_channels: int, out_channels: int, kernel_size: int, num_rotations: int = 4, stride: int = 1, padding: int = 0, bias: bool = True, device: str = 'cuda')[source]

Bases: Module

This class represents a roto-reflection equivariant convolutional layer with lifting.

The layer is equivariant to a group of rotations and reflections. The weights of the layer are initialized using the Kaiming uniform initialization method. The layer supports optional bias.

__init__()[source]

Initializes the RotoReflectionEquivariantConvLift instance.

get_rotoreflected_weights()[source]

Returns the weights of the layer after rotation, reflection, and permutation.

forward()[source]

Performs a forward pass through the layer.

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the layer.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the layer. It has the shape (batch_size, out_channels, num_group_elements, height, width).

Return type:

torch.Tensor

get_rotoreflected_weights(weights: Tensor, num_rotations: int = 4) → Tensor[source]

Returns the weights of the layer after rotation and reflection.

Parameters:

  • weights (torch.Tensor) – The weights of the layer.

  • num_rotations (int, optional) – The number of rotations in the group. Defaults to 4.

Returns:

The weights after rotation, reflection, and permutation.

Return type:

torch.Tensor

class equiadapt.images.SteerableImageCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig, in_shape: tuple)[source]

Bases: ContinuousGroupImageCanonicalization

This class represents a steerable image canonicalization model.

The model is designed to be equivariant under a continuous group of euclidean transformations - rotations and reflections.

__init__()[source]

Initializes the SteerableImageCanonicalization instance.

get_rotation_matrix_from_vector()[source]

This method takes the input vector and returns the rotation matrix.

get_groupelement()[source]

This method maps the input image to the group element.

get_groupelement(x: Tensor) → dict[source]

This method takes the input image and maps it to the group element

Parameters:

x (torch.Tensor) – input image

Returns:

group element

Return type:

dict

get_rotation_matrix_from_vector(vectors: Tensor) → Tensor[source]

This method takes the input vector and returns the rotation matrix

Parameters:

vectors (torch.Tensor) – input vector

Returns:

rotation matrices

Return type:

torch.Tensor

class equiadapt.images.WideResNet101Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)[source]

Bases: Module

This class represents a neural network based on the WideResNetNetwork architecture.

The network uses a pre-trained WideResNet model. The final fully connected layer of the WideResNet101 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

class equiadapt.images.WideResNet50Network(in_shape: tuple, out_channels: int, kernel_size: int, num_layers: int = 2, out_vector_size: int = 128)[source]

Bases: Module

This class represents a neural network based on the WideResNetNetwork architecture.

The network uses a pre-trained WideResNet model. The final fully connected layer of the WideResNet50 model is replaced with a new fully connected layer.

resnet18

The ResNet-18 model.

Type:

torchvision.models.ResNet

out_vector_size

The size of the output vector of the network.

Type:

int

forward(x: Tensor) → Tensor[source]

Performs a forward pass through the network.

Parameters:

x (torch.Tensor) – The input data. It should have the shape (batch_size, in_channels, height, width).

Returns:

The output of the network. It has the shape (batch_size, 1).

Return type:

torch.Tensor

equiadapt.images.flip_boxes(boxes: Tensor, width: int) → Tensor[source]

Flips bounding boxes horizontally.

Parameters:

  • boxes (torch.Tensor) – The bounding boxes to flip.

  • width (int) – The width of the image.

Returns:

The flipped bounding boxes.

Return type:

torch.Tensor

equiadapt.images.flip_masks(masks: Tensor) → Tensor[source]

Flips masks horizontally.

Parameters:

masks (torch.Tensor) – The masks to flip.

Returns:

The flipped masks.

Return type:

torch.Tensor

equiadapt.images.get_action_on_image_features(feature_map: Tensor, group_info_dict: dict, group_element_dict: dict, induced_rep_type: str = 'regular') → Tensor[source]

Applies a group action to the feature map.

Parameters:

  • feature_map (torch.Tensor) – The input feature map.

  • group_info_dict (dict) – A dictionary containing information about the group.

  • group_element_dict (dict) – A dictionary containing the group elements.

  • induced_rep_type (str, optional) – The type of induced representation. Defaults to “regular”.

Returns:

The feature map after the group action has been applied.

Return type:

torch.Tensor

equiadapt.images.roll_by_gather(feature_map: Tensor, shifts: Tensor) → Tensor[source]

Shifts the feature map along the group dimension by the specified shifts.

Parameters:

  • feature_map (torch.Tensor) – The input feature map. It should have the shape (batch, channel, group, x_dim, y_dim).

  • shifts (torch.Tensor) – The shifts for each feature map in the batch.

Returns:

The shifted feature map.

Return type:

torch.Tensor

equiadapt.images.rotate_boxes(boxes: Tensor, angle: Tensor, width: int) → Tensor[source]

Rotates bounding boxes by a specified angle.

Parameters:

  • boxes (torch.Tensor) – The bounding boxes to rotate.

  • angle (torch.Tensor) – The angle to rotate the bounding boxes by.

  • width (int) – The width of the image.

Returns:

The rotated bounding boxes.

Return type:

torch.Tensor

equiadapt.images.rotate_masks(masks: Tensor, angle: Tensor) → Tensor[source]

Rotates masks by a specified angle.

Parameters:

  • masks (torch.Tensor) – The masks to rotate.

  • angle (torch.Tensor) – The angle to rotate the masks by.

Returns:

The rotated masks.

Return type:

torch.Tensor

equiadapt.images.rotate_points(origin: List[float], point: Tensor, angle: Tensor) → Tuple[Tensor, Tensor][source]

Rotates a point around an origin by a specified angle.

Parameters:

  • origin (List[float]) – The origin to rotate the point around.

  • point (torch.Tensor) – The point to rotate.

  • angle (torch.Tensor) – The angle to rotate the point by.

Returns:

The rotated point.

Return type:

Tuple[torch.Tensor, torch.Tensor]