transforms

class mfai.pytorch.transforms.DimensionSubSampler(dim_name, idx_to_keep)[source]

Bases: Module

Subsamples a NamedTensor dimension.

Exemple: ``` import torch from mfai.pytorch.namedtensor import NamedTensor

data = NamedTensor(

tensor=torch.rand(10, 512, 512, 3), names=[“timesteps”, “lat”, “lon”, “features”], feature_names=[“u”, “v”, “t2m”],

)

sub_sampler = DimensionSubSampler(

dim_name=”timesteps”, idx_to_keep=[0, 3, 6, 9],

)

transformed_data = sub_sampler(data)

print(“data:”) print(data) print(“transformed_data:”) print(transformed_data) ` => ` data: — NamedTensor — Names: [‘timesteps’, ‘lat’, ‘lon’, ‘features’] Tensor Shape: torch.Size([10, 512, 512, 3])) transformed_data: — NamedTensor — Names: [‘timesteps’, ‘lat’, ‘lon’, ‘features’] Tensor Shape: torch.Size([4, 512, 512, 3])) ```

Parameters:
forward(named_tensor)[source]

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: NamedTensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

named_tensor (NamedTensor)

Return type:

NamedTensor

class mfai.pytorch.transforms.MeanDimensionSubsampler(dim_name, idx_to_be_meaned)[source]

Bases: Module

Subsamples a NamedTensor dimension by averaging the specified indexes.

Exemple: ``` import torch from mfai.pytorch.namedtensor import NamedTensor

data = NamedTensor(

tensor=torch.rand(10, 512, 512, 3), names=[“timesteps”, “lat”, “lon”, “features”], feature_names=[“u”, “v”, “t2m”],

)

sub_sampler = DimensionSubSampler(

dim_name=”timesteps”, idx_to_be_meaned=[

[0, 1, 2], [2, 3, 4], [4, 5, 6], [6, 7, 8],

],

)

transformed_data = sub_sampler(data)

print(“data:”) print(data) print(“transformed_data:”) print(transformed_data) ` => ` data: — NamedTensor — Names: [‘timesteps’, ‘lat’, ‘lon’, ‘features’] Tensor Shape: torch.Size([10, 512, 512, 3])) transformed_data: — NamedTensor — Names: [‘timesteps’, ‘lat’, ‘lon’, ‘features’] Tensor Shape: torch.Size([4, 512, 512, 3])) ```

Parameters:
forward(named_tensor)[source]

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: NamedTensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

named_tensor (NamedTensor)

Return type:

NamedTensor

class mfai.pytorch.transforms.RandomCropWithMinPositivePixels(crop_size=(512, 512), min_positive_percentage=15.0, tries=5)[source]

Bases: object

Randomly crops an input image to a 512x512 image, with min 15% of positive pixels. The random crop is performed 5 times at the most, or until an image with 15% of positive pixels is found. If such an image is not found, it returns the image with the most positive pixels from the 5 images cropped.

input is a tuple sample with x and y x shape (C, H, W) y shape (1, H, W)

Parameters:
get_crop(y)[source]
Return type:

tuple[int, int]

Parameters:

y (Tensor)

get_one_crop(y)[source]
Return type:

tuple[int, int, float]

Parameters:

y (Tensor)