Environments module¶

Environments¶

activemri.envs.envs.py¶

Gym-like environment for active MRI acquisition.

class activemri.envs.envs.ActiveMRIEnv(kspace_shape: Tuple[int, int], num_parallel_episodes: int = 1, budget: Optional[int] = None, seed: Optional[int] = None)¶

Bases: gym.core.Env

Base class for all active MRI acquisition environments.

This class provides the core logic implementation of the k-space acquisition process. The class is not to be used directly, but rather one of its subclasses should be instantiated. Subclasses of ActiveMRIEnv are responsible for data initialization and specifying configuration options for the environment.

Parameters

kspace_shape (tuple(int,int)) – Shape of the k-space slices for the dataset.
num_parallel_episodes (int) – Determines the number images that will be processed simultaneously by reset() and step(). Defaults to 1.
budget (optional(int)) – The length of an acquisition episode. Defaults to None, which indicates that episode will continue until all k-space columns have been acquired.
seed (optional(int)) – The seed for the environment’s random number generator, which is an instance of numpy.random.RandomState. Defaults to None.
no_checkpoint (optional(bool)) – Set to True if you want to run your reconstructor model without loading anything from a checkpoint.

render(mode='human')¶

Renders information about the environment’s current state.

Returns

An image frame containing, from left to right: current: acquisition mask, current ground image, current reconstruction, and current relative reconstruction error.

Return type

np.ndarray

reset() → Tuple[Dict[str, Any], Dict[str, Any]]¶

Starts a new acquisition episode with a batch of images.

This methods performs the following steps:

Reads a batch of images from the environment’s dataset.

Creates an initial acquisition mask for each image.

Passes the loaded data and the initial masks to the transform function, producing a batch of inputs for the environment’s reconstructor model.

Calls the reconstructor model on this input and returns its output as an observation.

The observation returned is a dictionary with the following keys:

“reconstruction”(torch.Tensor): The reconstruction produced by the environment’s reconstructor model, using the current acquisition mask.
“extra_outputs”(dict(str,Any)): A dictionary with any additional outputs produced by the reconstructor (e.g., uncertainty maps).
“mask”(torch.Tensor): The current acquisition mask.

Returns

tuple containing:

obs(dict(str,any): Observation dictionary.
metadata(dict(str,any): Metadata information containing the following keys:
- ”fname”(list(str)): the filenames of the image read from the dataset.
- ”slice_id”(list(int)): slice indices for each image within the volume.
- ”current_score”(dict(str,float): A dictionary with the error measures for the reconstruction (e.g., “mse”, “nmse”, “ssim”, “psnr”). The measures considered can be obtained with score_keys().

Return type

tuple

static score_keys() → List[str]¶: Returns the list of score metric names used by this environment.

seed(seed: Optional[int] = None)¶

Sets the seed for the internal number generator.

This seeds affects the order of the data loader for all loop modalities (i.e., training, validation, test).

Parameters: seed (optional(int)) – The seed for the environment’s random number generator.

set_test(reset: bool = True)¶

Sets the environment to use the test data loader.

Parameters: reset (bool) – If True, also resets the data loader so that it starts again from the first image in the loop order.

Warning

After this method is called the env.reset() needs to be called again, otherwise an exception will be thrown.

set_training(reset: bool = False)¶

Sets the environment to use the training data loader.

Parameters: reset (bool) – If True, also resets the data loader so that it starts again from the first image in the loop order.

Warning

After this method is called the env.reset() needs to be called again, otherwise an exception will be thrown.

set_val(reset: bool = True)¶

Sets the environment to use the validation data loader.

Parameters: reset (bool) – If True, also resets the data loader so that it starts again from the first image in the loop order.

Warning

After this method is called the env.reset() needs to be called again, otherwise an exception will be thrown.

step(action: Union[int, Sequence[int]]) → Tuple[Dict[str, Any], numpy.ndarray, List[bool], Dict]¶

Performs a step of active MRI acquisition.

Given a set of indices for k-space columns to acquire, updates the current batch of masks with their corresponding indices, creates a new batch of reconstructions, and returns the corresponding observations and rewards (for the observation format see reset()). The reward is the improvement in score with respect to the reconstruction before adding the indices. The specific score metric used is determined by env.reward_metric.

The method also returns a list of booleans, indicating whether any episodes in the batch have already concluded.

The last return value is a metadata dictionary. It contains a single key “current_score”, which contains a dictionary with the error measures for the reconstruction (e.g., "mse", "nmse", "ssim", "psnr"). The measures considered can be obtained with score_keys().

Parameters

action (union(int, sequence(int))) – Indices for k-space columns to acquire. The length of the sequence must be equal to the current number of parallel episodes (i.e., obs["reconstruction"].shape[0]). If only an int is passed, the index will be replicated for the whole batch of episodes.

Returns

The transition information in the order (next_observation, reward, done, meta). The types and shapes are:

next_observation(dict): Dictionary format (see reset()).

reward(np.ndarray): length equal to current number of parallel episodes.

done(list(bool)): same length as reward.

meta(dict): see description above.

Return type

tuple

try_action(action: Union[int, Sequence[int]]) → Tuple[Dict[str, Any], Dict[str, numpy.ndarray]]¶

Simulates the effects of actions without changing the environment’s state.

This method operates almost exactly as step(), with the exception that the environment’s state is not altered. The method returns the next observation and the resulting reconstruction score after applying the give k-space columns to each image in the current batch of episodes.

Parameters

action (union(int, sequence(int))) – Indices for k-space columns to acquire. The length of the sequence must be equal to the current number of parallel episodes (i.e., obs["reconstruction"].shape[0]). If only an int is passed, the index will be replicated for the whole batch of episodes.

Returns

The reconstruction information in the order (next_observation, current_score). The types and shapes are:

next_observation(dict): Dictionary format (see reset()).

current_score(dict(str, float)): A dictionary with the error measures
for the reconstruction (e.g., “mse”, “nmse”, “ssim”, “psnr”). The measures considered can be obtained with ActiveMRIEnv.score_keys().

Return type

tuple

class activemri.envs.envs.CyclicSampler(data_source: Sized, order: Optional[Sized] = None, loops: int = 1)¶: Bases: torch.utils.data.sampler.Sampler

class activemri.envs.envs.FastMRIEnv(config_path: str, dataset_name: str, num_parallel_episodes: int = 1, budget: Optional[int] = None, seed: Optional[int] = None, num_cols: Sequence[int] = 368, 372)¶

Bases: activemri.envs.envs.ActiveMRIEnv

Base class for all fastMRI environments.

This class can be used to instantiate active acquisition environments using fastMRI data. However, for convenience we provide subclasses of FastMRIEnv with default configuration options for each dataset:

SingleCoilKneeEnv

MultiCoilKneeEnv

SingleCoilBrainEnv

MultiCoilKneeEnv

Parameters

config_path (str) – The path to the JSON configuration file.
dataset_name (str) – One of “knee_singlecoil”, “multicoil” (for knee), “brain_multicoil”. Primarily used to locate the fastMRI dataset in the user’s fastMRI data root folder.
num_parallel_episodes (int) – Determines the number images that will be processed simultaneously by reset() and step(). Defaults to 1.
budget (optional(int)) – The length of an acquisition episode. Defaults to None, which indicates that episode will continue until all k-space columns have been acquired.
seed (optional(int)) – The seed for the environment’s random number generator, which is an instance of numpy.random.RandomState. Defaults to None.
num_cols (sequence(int)) – Used to filter k-space data to only use images whose k-space width is in this tuple. Defaults to (368, 372).

render(mode='human')¶

Renders information about the environment’s current state.

Returns

An image frame containing, from left to right: current: acquisition mask, current ground image, current reconstruction, and current relative reconstruction error.

Return type

np.ndarray

class activemri.envs.envs.MICCAI2020Env(num_parallel_episodes: int = 1, budget: Optional[int] = None, seed: Optional[int] = None, extreme: bool = False, obs_includes_padding: bool = True)¶

Bases: activemri.envs.envs.ActiveMRIEnv

Implementation of environment used for Pineda et al., MICCAI 2020.

This environment is provided to facilitate replication of the experiments performed in Luis Pineda, Sumana Basu, Adriana Romero, Roberto Calandra, Michal Drozdzal, “Active MR k-space Sampling with Reinforcement Learning”. MICCAI 2020.

The dataset is the same as that of SingleCoilKneeEnv, except that we provide a custom validation/test split of the original validation data. The environment’s configuration file is set to use the reconstruction model used in the paper (see activemri.models.cvpr19_reconstructor.CVPR19Reconstructor), as well as the proper transform to generate inputs for this model.

The k-space shape of this environment is set to (640, 368).

Parameters

num_parallel_episodes (int) – Determines the number images that will be processed simultaneously by reset() and step(). Defaults to 1.
budget (optional(int)) – The length of an acquisition episode. Defaults to None, which indicates that episode will continue until all k-space columns have been acquired.
seed (optional(int)) – The seed for the environment’s random number generator, which is an instance of numpy.random.RandomState. Defaults to None.
extreme (bool) – True or False for running extreme acceleration or normal acceleration scenarios described in the paper, respectively.

render(mode='human')¶

Renders information about the environment’s current state.

Returns

An image frame containing, from left to right: current: acquisition mask, current ground image, current reconstruction, and current relative reconstruction error.

Return type

np.ndarray

reset() → Tuple[Dict[str, Any], Dict[str, Any]]¶

Starts a new acquisition episode with a batch of images.

This methods performs the following steps:

Reads a batch of images from the environment’s dataset.

Creates an initial acquisition mask for each image.

Passes the loaded data and the initial masks to the transform function, producing a batch of inputs for the environment’s reconstructor model.

Calls the reconstructor model on this input and returns its output as an observation.

The observation returned is a dictionary with the following keys:

“reconstruction”(torch.Tensor): The reconstruction produced by the environment’s reconstructor model, using the current acquisition mask.
“extra_outputs”(dict(str,Any)): A dictionary with any additional outputs produced by the reconstructor (e.g., uncertainty maps).
“mask”(torch.Tensor): The current acquisition mask.

Returns

tuple containing:

obs(dict(str,any): Observation dictionary.
metadata(dict(str,any): Metadata information containing the following keys:
- ”fname”(list(str)): the filenames of the image read from the dataset.
- ”slice_id”(list(int)): slice indices for each image within the volume.
- ”current_score”(dict(str,float): A dictionary with the error measures for the reconstruction (e.g., “mse”, “nmse”, “ssim”, “psnr”). The measures considered can be obtained with score_keys().

Return type

tuple

step(action: Union[int, Sequence[int]]) → Tuple[Dict[str, Any], numpy.ndarray, List[bool], Dict]¶

Performs a step of active MRI acquisition.

Given a set of indices for k-space columns to acquire, updates the current batch of masks with their corresponding indices, creates a new batch of reconstructions, and returns the corresponding observations and rewards (for the observation format see reset()). The reward is the improvement in score with respect to the reconstruction before adding the indices. The specific score metric used is determined by env.reward_metric.

The method also returns a list of booleans, indicating whether any episodes in the batch have already concluded.

The last return value is a metadata dictionary. It contains a single key “current_score”, which contains a dictionary with the error measures for the reconstruction (e.g., "mse", "nmse", "ssim", "psnr"). The measures considered can be obtained with score_keys().

Parameters

action (union(int, sequence(int))) – Indices for k-space columns to acquire. The length of the sequence must be equal to the current number of parallel episodes (i.e., obs["reconstruction"].shape[0]). If only an int is passed, the index will be replicated for the whole batch of episodes.

Returns

The transition information in the order (next_observation, reward, done, meta). The types and shapes are:

next_observation(dict): Dictionary format (see reset()).

reward(np.ndarray): length equal to current number of parallel episodes.

done(list(bool)): same length as reward.

meta(dict): see description above.

Return type

tuple

class activemri.envs.envs.MultiCoilKneeEnv(num_parallel_episodes: int = 1, budget: Optional[int] = None, seed: Optional[int] = None, num_cols: Sequence[int] = 368, 372)¶

Bases: activemri.envs.envs.FastMRIEnv

Convenience class to access multi-coil knee data.

Loads the configuration from configs/multi-coil-knee.json. Looks for datasets named “multicoil_{train/val/test}” under default data_location dir. If “test” is not found, it uses “val” folder for test mode.

Parameters

num_parallel_episodes (int) – Determines the number images that will be processed simultaneously by reset() and step(). Defaults to 1.
budget (optional(int)) – The length of an acquisition episode. Defaults to None, which indicates that episode will continue until all k-space columns have been acquired.
seed (optional(int)) – The seed for the environment’s random number generator, which is an instance of numpy.random.RandomState. Defaults to None.
num_cols (sequence(int)) – Used to filter k-space data to only use images whose k-space width is in this tuple. Defaults to (368, 372).

class activemri.envs.envs.SingleCoilKneeEnv(num_parallel_episodes: int = 1, budget: Optional[int] = None, seed: Optional[int] = None, num_cols: Sequence[int] = 368, 372)¶

Bases: activemri.envs.envs.FastMRIEnv

Convenience class to access single-coil knee data.

Loads the configuration from configs/single-coil-knee.json. Looks for datasets named “knee_singlecoil_{train/val/test}” under the data_location dir. If “test” is not found, it uses “val” folder for test mode.

Parameters

num_parallel_episodes (int) – Determines the number images that will be processed simultaneously by reset() and step(). Defaults to 1.
budget (optional(int)) – The length of an acquisition episode. Defaults to None, which indicates that episode will continue until all k-space columns have been acquired.
seed (optional(int)) – The seed for the environment’s random number generator, which is an instance of numpy.random.RandomState. Defaults to None.
num_cols (sequence(int)) – Used to filter k-space data to only use images whose k-space width is in this tuple. Defaults to (368, 372).

Mask utilities¶

activemri.envs.masks.py¶

Utilities to generate and manipulate active acquisition masks.

activemri.envs.masks.sample_low_frequency_mask(mask_args: Dict[str, Any], kspace_shapes: List[Tuple[int, …]], rng: numpy.random.mtrand.RandomState, attrs: Optional[List[Dict[str, Any]]] = None) → torch.Tensor¶

Samples low frequency masks.

Returns masks that contain some number of the lowest k-space frequencies active. The number of frequencies doesn’t have to be the same for all masks in the batch, and it can also be a random number, depending on the given mask_args. Active columns will be represented as 1s in the mask, and inactive columns as 0s.

The distribution and shape of the masks can be controlled by mask_args. This is a dictionary with the following keys:

“max_width”(int): The maximum width of the masks.

“min_cols”(int): The minimum number of low frequencies columns to activate per side.

“max_cols”(int): The maximum number of low frequencies columns to activate per side (inclusive).

“width_dim”(int): Indicates which of the dimensions in kspace_shapes corresponds to the k-space width.

“centered”(bool): Specifies if the low frequencies are in the center of the k-space (True) or on the edges (False).

“apply_attrs_padding”(optional(bool)): If True, the function will read keys "padding_left" and "padding_right" from attrs and set all corresponding high-frequency columns to 1.

The number of 1s in the effective region of the mask (see next paragraph) is sampled between mask_args["min_cols"] and mask_args["max_cols"] (inclusive). The number of dimensions for the mask tensor will be mask_args["width_dim"] + 2. The size will be [batch_size, 1, ..., 1, mask_args["max_width"]]. For example, with mask_args["width_dim"] = 1 and mask_args["max_width"] = 368, output tensor has shape [batch_size, 1, 368].

This function supports simultaneously sampling masks for k-space of different number of columns. This is controlled by argument kspace_shapes. From this list, the function will obtain 1) batch_size = len(kspace_shapes), and 2) the width of the k-spaces for each element in the batch. The i-th mask will have kspace_shapes[item][mask_args["width_dim"]] effective columns.

Note

The mask tensor returned will always have mask_args["max_width"] columns. However, for any element i s.t. kspace_shapes[i][mask_args["width_dim"]] < mask_args["max_width"], the function will then pad the extra k-space columns with 1s. The rest of the columns will be filled out as if the mask has the same width as that indicated by kspace_shape[i].

Parameters

mask_args (dict(str,any)) – Specifies configuration options for the masks, as explained above.
kspace_shapes (list(tuple(int,..))) – Specifies the shapes of the k-space data on which this mask will be applied, as explained above.
rng (np.random.RandomState) – A random number generator to sample the masks.
attrs (dict(str,int)) – Used to determine any high-frequency padding. It must contain keys "padding_left" and "padding_right".

Returns

The generated low frequency masks.

Return type

torch.Tensor