Shortcuts

mmedit.models.utils

Package Contents

Functions

extract_around_bbox(img, bbox, target_size[, ...])

Extract patches around the given bbox.

extract_bbox_patch(bbox, img[, channel_first])

Extract patch from a given bbox.

flow_warp(x, flow[, interpolation, padding_mode, ...])

Warp an image or a feature map with optical flow.

default_init_weights(module[, scale])

Initialize network weights.

generation_init_weights(module[, init_type, init_gain])

Default initialization of network weights for image generation.

get_module_device(module)

Get the device of a module.

get_valid_noise_size(→ Optional[int])

Get the value of noise_size from input, generator and check the

get_valid_num_batches(→ int)

Try get the valid batch size from inputs.

make_layer(block, num_blocks, **kwarg)

Make layers by stacking the same blocks.

set_requires_grad(nets[, requires_grad])

Set requires_grad for all the networks.

label_sample_fn(→ Union[torch.Tensor, None])

Sample random label with respect to num_batches, num_classes and

noise_sample_fn(→ torch.Tensor)

Sample noise with respect to the given num_batches, noise_size and

get_unknown_tensor(trimap[, unknown_value])

Get 1-channel unknown area tensor from the 3 or 1-channel trimap tensor.

normalize_vecs(→ torch.Tensor)

Normalize vector with it's lengths at the last dimension. If vector is

mmedit.models.utils.extract_around_bbox(img, bbox, target_size, channel_first=True)[source]

Extract patches around the given bbox.

Parameters
  • img (torch.Tensor | numpy.array) – Image data to be extracted. If organized in batch dimension, the batch dimension must be the first order like (n, h, w, c) or (n, c, h, w).

  • bbox (np.ndarray | torch.Tensor) – Bboxes to be modified. Bbox can be in batch or not.

  • target_size (List(int)) – Target size of final bbox.

  • channel_first (bool) – If True, the channel dimension of img is before height and width, e.g. (c, h, w). Otherwise, the img shape (samples in the batch) is like (h, w, c). Default: True.

Returns

Extracted patches. The dimension of the output should be the same as img.

Return type

(torch.Tensor | np.ndarray)

mmedit.models.utils.extract_bbox_patch(bbox, img, channel_first=True)[source]

Extract patch from a given bbox.

Parameters
  • bbox (torch.Tensor | numpy.array) – Bbox with (top, left, h, w). If img has batch dimension, the bbox must be stacked at first dimension. The shape should be (4,) or (n, 4).

  • img (torch.Tensor | numpy.array) – Image data to be extracted. If organized in batch dimension, the batch dimension must be the first order like (n, h, w, c) or (n, c, h, w).

  • channel_first (bool) – If True, the channel dimension of img is before height and width, e.g. (c, h, w). Otherwise, the img shape (samples in the batch) is like (h, w, c). Default: True.

Returns

Extracted patches. The dimension of the output should be the same as img.

Return type

(torch.Tensor | numpy.array)

mmedit.models.utils.flow_warp(x, flow, interpolation='bilinear', padding_mode='zeros', align_corners=True)[source]

Warp an image or a feature map with optical flow.

Parameters
  • x (Tensor) – Tensor with size (n, c, h, w).

  • flow (Tensor) – Tensor with size (n, h, w, 2). The last dimension is a two-channel, denoting the width and height relative offsets. Note that the values are not normalized to [-1, 1].

  • interpolation (str) – Interpolation mode: ‘nearest’ or ‘bilinear’. Default: ‘bilinear’.

  • padding_mode (str) – Padding mode: ‘zeros’ or ‘border’ or ‘reflection’. Default: ‘zeros’.

  • align_corners (bool) – Whether align corners. Default: True.

Returns

Warped image or feature map.

Return type

Tensor

mmedit.models.utils.default_init_weights(module, scale=1)[source]

Initialize network weights.

Parameters
  • modules (nn.Module) – Modules to be initialized.

  • scale (float) – Scale initialized weights, especially for residual blocks. Default: 1.

mmedit.models.utils.generation_init_weights(module, init_type='normal', init_gain=0.02)[source]

Default initialization of network weights for image generation.

By default, we use normal init, but xavier and kaiming might work better for some applications.

Parameters
  • module (nn.Module) – Module to be initialized.

  • init_type (str) – The name of an initialization method: normal | xavier | kaiming | orthogonal. Default: ‘normal’.

  • init_gain (float) – Scaling factor for normal, xavier and orthogonal. Default: 0.02.

mmedit.models.utils.get_module_device(module)[source]

Get the device of a module.

Parameters

module (nn.Module) – A module contains the parameters.

Returns

The device of the module.

Return type

torch.device

mmedit.models.utils.get_valid_noise_size(noise_size: Optional[int], generator: Union[Dict, torch.nn.Module]) Optional[int][source]

Get the value of noise_size from input, generator and check the consistency of these values. If no conflict is found, return that value.

Parameters
  • noise_size (Optional[int]) – noise_size passed to BaseGAN_refactor’s initialize function.

  • generator (ModelType) – The config or the model of generator.

Returns

The noise size feed to generator.

Return type

int | None

mmedit.models.utils.get_valid_num_batches(batch_inputs: mmedit.utils.typing.ForwardInputs) int[source]

Try get the valid batch size from inputs.

  • If some values in batch_inputs are Tensor and ‘num_batches’ is in batch_inputs, we check whether the value of ‘num_batches’ and the the length of first dimension of all tensors are same. If the values are not same, AssertionError will be raised. If all values are the same, return the value.

  • If no values in batch_inputs is Tensor, ‘num_batches’ must be contained in batch_inputs. And this value will be returned.

  • If some values are Tensor and ‘num_batches’ is not contained in batch_inputs, we check whether all tensor have the same length on the first dimension. If the length are not same, AssertionError will be raised. If all length are the same, return the length as batch size.

  • If batch_inputs is a Tensor, directly return the length of the first dimension as batch size.

Parameters

batch_inputs (ForwardInputs) – Inputs passed to forward().

Returns

The batch size of samples to generate.

Return type

int

mmedit.models.utils.make_layer(block, num_blocks, **kwarg)[source]

Make layers by stacking the same blocks.

Parameters
  • block (nn.module) – nn.module class for basic block.

  • num_blocks (int) – number of blocks.

Returns

Stacked blocks in nn.Sequential.

Return type

nn.Sequential

mmedit.models.utils.set_requires_grad(nets, requires_grad=False)[source]

Set requires_grad for all the networks.

Parameters
  • nets (nn.Module | list[nn.Module]) – A list of networks or a single network.

  • requires_grad (bool) – Whether the networks require gradients or not

mmedit.models.utils.label_sample_fn(label: Union[torch.Tensor, Callable, List[int], None] = None, *, num_batches: int = 1, num_classes: Optional[int] = None, device: Optional[str] = None) Union[torch.Tensor, None][source]

Sample random label with respect to num_batches, num_classes and device.

Parameters
  • label (Union[Tensor, Callable, List[int], None], optional) – You can directly give a batch of label through a torch.Tensor or offer a callable function to sample a batch of label data. Otherwise, the None indicates to use the default label sampler. Defaults to None.

  • num_batches (int, optional) – The number of batch size. Defaults to 1.

  • num_classes (Optional[int], optional) – The number of classes. Defaults to None.

  • device (Optional[str], optional) – The target device of the label. Defaults to None.

Returns

Sampled random label.

Return type

Union[Tensor, None]

mmedit.models.utils.noise_sample_fn(noise: Union[torch.Tensor, Callable, None] = None, *, num_batches: int = 1, noise_size: Union[int, Sequence[int], None] = None, device: Optional[str] = None) torch.Tensor[source]

Sample noise with respect to the given num_batches, noise_size and device.

Parameters
  • noise (torch.Tensor | callable | None) – You can directly give a batch of noise through a torch.Tensor or offer a callable function to sample a batch of noise data. Otherwise, the None indicates to use the default noise sampler. Defaults to None.

  • num_batches (int, optional) – The number of batch size. Defaults to 1.

  • noise_size (Union[int, Sequence[int], None], optional) – The size of random noise. Defaults to None.

  • device (Optional[str], optional) – The target device of the random noise. Defaults to None.

Returns

Sampled random noise.

Return type

Tensor

mmedit.models.utils.get_unknown_tensor(trimap, unknown_value=128 / 255)[source]

Get 1-channel unknown area tensor from the 3 or 1-channel trimap tensor.

Parameters
  • trimap (Tensor) – Tensor with shape (N, 3, H, W) or (N, 1, H, W).

  • unknown_value (float) – Scalar value indicating unknown region in trimap. If trimap is pre-processed using ‘rescale_to_zero_one’, then 0 for bg, 128/255 for unknown, 1 for fg, and unknown_value should set to 128 / 255. If trimap is pre-processed by FormatTrimap(to_onehot=False)(), then 0 for bg, 1 for unknown, 2 for fg and unknown_value should set to 1. If trimap is pre-processed by FormatTrimap(to_onehot=True)(), then trimap is 3-channeled, and this value is not used.

Returns

Unknown area mask of shape (N, 1, H, W).

Return type

Tensor

mmedit.models.utils.normalize_vecs(vectors: torch.Tensor) torch.Tensor[source]

Normalize vector with it’s lengths at the last dimension. If vector is two-dimension tensor, this function is same as L2 normalization.

Parameters

vector (torch.Tensor) – Vectors to be normalized.

Returns

Vectors after normalization.

Return type

torch.Tensor

Read the Docs v: latest
Versions
master
latest
stable
zyh-re-docs
zyh-doc-notfound-extend
zyh-api-rendering
Downloads
pdf
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.