diffengine.datasets.transforms.loading

Module Contents

Classes

LoadMask

Load Mask for multiple types.

Functions

random_bbox(img_shape, max_bbox_shape[, ...])

Generate a random bbox for the mask on a given image.

bbox2mask(img_shape, bbox[, dtype])

Generate mask in np.ndarray from bbox.

random_irregular_mask(img_shape[, num_vertices, ...])

Generate random irregular masks.

get_irregular_mask(img_shape[, area_ratio_range])

Get irregular mask with the constraints in mask ratio.

brush_stroke_mask(img_shape[, num_vertices, ...])

Generate free-form mask.
diffengine.datasets.transforms.loading.random_bbox(img_shape, max_bbox_shape, max_bbox_delta=40, min_margin=20)[source]

Generate a random bbox for the mask on a given image.

Copied from https://github.com/open-mmlab/mmagic/blob/main/mmagic/utils/trans_utils.py

In our implementation, the max value cannot be obtained since we use np.random.randint. And this may be different with other standard scripts in the community.

Args:

img_shape (tuple[int]): The size of a image, in the form of (h, w). max_bbox_shape (int | tuple[int]): Maximum shape of the mask box,

in the form of (h, w). If it is an integer, the mask box will be square.

max_bbox_delta (int | tuple[int]): Maximum delta of the mask box,

in the form of (delta_h, delta_w). If it is an integer, delta_h and delta_w will be the same. Mask shape will be randomly sampled from the range of max_bbox_shape - max_bbox_delta and max_bbox_shape. Default: (40, 40).

min_margin (int | tuple[int]): The minimum margin size from the

edges of mask box to the image boarder, in the form of (margin_h, margin_w). If it is an integer, margin_h and margin_w will be the same. Default: (20, 20).

Returns:

tuple[int]: The generated box, (top, left, h, w).

Parameters:

  • img_shape (tuple[int, int]) –

  • max_bbox_shape (int | tuple[int, int]) –

  • max_bbox_delta (int | tuple[int, int]) –

  • min_margin (int | tuple[int, int]) –

Return type:

tuple[int, int, int, int]

diffengine.datasets.transforms.loading.bbox2mask(img_shape, bbox, dtype='uint8')[source]

Generate mask in np.ndarray from bbox.

Copied from https://github.com/open-mmlab/mmagic/blob/main/mmagic/utils/trans_utils.py

The returned mask has the shape of (h, w, 1). ‘1’ indicates the hole and ‘0’ indicates the valid regions.

We prefer to use uint8 as the data type of masks, which may be different from other codes in the community.

Args:

img_shape (tuple[int]): The size of the image. bbox (tuple[int]): Configuration tuple, (top, left, height, width) np.dtype (str): Indicate the data type of returned masks.

Default: ‘uint8’

Returns:

mask (np.ndarray): Mask in the shape of (h, w, 1).

Parameters:

  • img_shape (tuple[int, int]) –

  • bbox (tuple[int, int, int, int]) –

  • dtype (str) –

Return type:

numpy.ndarray

diffengine.datasets.transforms.loading.random_irregular_mask(img_shape, num_vertices=(4, 8), max_angle=4, length_range=(10, 100), brush_width=(10, 40), dtype='uint8')[source]

Generate random irregular masks.

Copied from https://github.com/open-mmlab/mmagic/blob/main/mmagic/utils/trans_utils.py

This is a modified version of free-form mask implemented in ‘brush_stroke_mask’.

We prefer to use uint8 as the data type of masks, which may be different from other codes in the community.

Args:

img_shape (tuple[int]): Size of the image. num_vertices (int | tuple[int]): Min and max number of vertices. If

only give an integer, we will fix the number of vertices. Default: (4, 8).

max_angle (float): Max value of angle at each vertex. Default 4.0. length_range (int | tuple[int]): (min_length, max_length). If only give

an integer, we will fix the length of brush. Default: (10, 100).

brush_width (int | tuple[int]): (min_width, max_width). If only give

an integer, we will fix the width of brush. Default: (10, 40).

np.dtype (str): Indicate the data type of returned masks.

Default: ‘uint8’

Returns:

mask (np.ndarray): Mask in the shape of (h, w, 1).

Parameters:

  • img_shape (tuple[int, int]) –

  • num_vertices (int | tuple[int, int]) –

  • max_angle (float) –

  • length_range (int | tuple[int, int]) –

  • brush_width (int | tuple[int, int]) –

  • dtype (str) –

Return type:

numpy.ndarray

diffengine.datasets.transforms.loading.get_irregular_mask(img_shape, area_ratio_range=(0.15, 0.5), **kwargs)[source]

Get irregular mask with the constraints in mask ratio.

Copied from https://github.com/open-mmlab/mmagic/blob/main/mmagic/utils/trans_utils.py

Args:

img_shape (tuple[int]): Size of the image. area_ratio_range (tuple(float)): Contain the minimum and maximum area ratio. Default: (0.15, 0.5).

Returns:

mask (np.ndarray): Mask in the shape of (h, w, 1).

Parameters:

  • img_shape (tuple[int, int]) –

  • area_ratio_range (tuple[float, float]) –

Return type:

numpy.ndarray

diffengine.datasets.transforms.loading.brush_stroke_mask(img_shape, num_vertices=(4, 12), mean_angle=2 * math.pi / 5, angle_range=2 * math.pi / 15, brush_width=(12, 40), max_loops=4, dtype='uint8')[source]

Generate free-form mask.

Copied from https://github.com/open-mmlab/mmagic/blob/main/mmagic/utils/trans_utils.py

The method of generating free-form mask is in the following paper: Free-Form Image Inpainting with Gated Convolution.

When you set the config of this type of mask. You may note the usage of np.random.randint and the range of np.random.randint is [left, right).

We prefer to use uint8 as the data type of masks, which may be different from other codes in the community.

Args:

img_shape (tuple[int]): Size of the image. num_vertices (int | tuple[int]): Min and max number of vertices. If

only give an integer, we will fix the number of vertices. Default: (4, 12).

mean_angle (float): Mean value of the angle in each vertex. The angle

is measured in radians. Default: 2 * math.pi / 5.

angle_range (float): Range of the random angle.

Default: 2 * math.pi / 15.

brush_width (int | tuple[int]): (min_width, max_width). If only give

an integer, we will fix the width of brush. Default: (12, 40).

max_loops (int): The max number of for loops of drawing strokes.

Default: 4.

np.dtype (str): Indicate the data type of returned masks.

Default: ‘uint8’.

Returns:

mask (np.ndarray): Mask in the shape of (h, w, 1).

Parameters:

  • img_shape (tuple[int, int]) –

  • num_vertices (int | tuple[int, int]) –

  • mean_angle (float) –

  • angle_range (float) –

  • brush_width (int | tuple[int, int]) –

  • max_loops (int) –

  • dtype (str) –

Return type:

numpy.ndarray

class diffengine.datasets.transforms.loading.LoadMask(mask_mode='bbox', mask_config=None)[source]

Bases: diffengine.datasets.transforms.base.BaseTransform

Load Mask for multiple types.

Copied from https://github.com/open-mmlab/mmagic/blob/main/mmagic/utils/trans_utils.py

Reference from: mmagic.datasets.transforms.loading.LoadMask

For different types of mask, users need to provide the corresponding config dict.

Example config for bbox:

config = dict(max_bbox_shape=128)

Example config for irregular:

config = dict(
    num_vertices=(4, 12),
    max_angle=4.,
    length_range=(10, 100),
    brush_width=(10, 40),
    area_ratio_range=(0.15, 0.5))

Example config for ff:

config = dict(
    num_vertices=(4, 12),
    mean_angle=1.2,
    angle_range=0.4,
    brush_width=(12, 40))

Args:

mask_mode (str): Mask mode in [‘bbox’, ‘irregular’, ‘ff’, ‘set’,

‘whole’]. Default: ‘bbox’. * bbox: square bounding box masks. * irregular: irregular holes. * ff: free-form holes from DeepFillv2. * set: randomly get a mask from a mask set. * whole: use the whole image as mask.

mask_config (dict): Params for creating masks. Each type of mask needs

different configs. Default: None.

transform(results)[source]

Transform function.

Args:
results (dict): A dict containing the necessary information and

data for augmentation.

Returns:

dict: A dict containing the processed data and information.

Parameters:

results (dict) –

Return type:

dict

Parameters:

  • mask_mode (str) –

  • mask_config (dict | None) –