utils.py
Various utility functions.
- class ptlflow.utils.utils.InputPadder(dims: Sequence[int], stride: int, size: Tuple[int, int] | None = None, two_side_pad: bool = True, pad_mode: str = 'replicate', pad_value: float = 0.0)[source]
Pads images such that dimensions are divisible by stride.
This is just a wrapper for ptlflow.utils.external.raft.InputPadder.
- class ptlflow.utils.utils.InputScaler(orig_shape: Tuple[int, int], stride: int | None = None, size: Tuple[int, int] | None = None, scale_factor: float | None = 1.0, interpolation_mode: str = 'bilinear', interpolation_align_corners: bool = False)[source]
Scale 2D torch.Tensor input to a target size, and then rescale it back to the original size.
- fill(x: Tensor, is_flow: bool = False) Tensor[source]
Scale the input to the target size specified during initialization.
- Parameters:
- xtorch.Tensor
The input to be scaled. Its shape must be (…, C, H, W), where … means any number of dimensions.
- is_flowbool
Whether the input is a flow field or not. If it is, then its values are multiplied by the rescale factor.
- Returns:
- torch.Tensor
The scaled input.
- unfill(x: Tensor, is_flow: bool = False) Tensor[source]
Scale the input to back to the original size defined during initialization.
- Parameters:
- xtorch.Tensor
The input to be rescaled back. Its shape must be (…, C, H, W), where … means any number of dimensions.
- is_flowbool
Whether the input is a flow field or not. If it is, then its values are multiplied by the rescale factor.
- Returns:
- torch.Tensor
The rescaled input.
- ptlflow.utils.utils.add_datasets_to_parser(parser: ArgumentParser, dataset_config_path: str) ArgumentParser[source]
Add dataset paths as parser arguments.
The dataset names and default paths are loaded from a yaml file.
- Parameters:
- parserArgumentParser
An initialized parser, this function will add more arguments to it.
- dataset_config_pathstr
The path to the yaml file containing the dataset paths.
- Returns:
- ArgumentParser
The updated parser.
- ptlflow.utils.utils.are_shapes_compatible(shape1: Sequence[int], shape2: Sequence[int]) bool[source]
Check if two tensor shapes are compatible.
Similar to PyTorch or Numpy, two shapes are considered “compatible” if either they have the same shape or if one shape can be broadcasted into the other. We consider two shapes compatible if, and only if: 1. their shapes have the same length (same number of dimension), and 2. each dimension size is either equal or at least one of them is one.
- Parameters:
- shape1Sequence[int]
The dimensions of the first shape.
- shape2Sequence[int]
The dimensions of the second shape.
- Returns:
- bool
Whether the two given shapes are compatible.
- ptlflow.utils.utils.bgr_val_as_tensor(bgr_val: float | Tuple[float, float, float] | ndarray | Tensor, reference_tensor: Tensor, bgr_tensor_shape_position: int = -3) Tensor[source]
Convert multiple types of BGR values given as input to a torch.Tensor where the BGR values are in the same position as a reference tensor.
The bgr values can be: - a single number, in which case it will be repeated three times to represent BGR. - a tuple, list, np.ndarray, or torch.Tensor with three elements.
The resulting tensor will have the BGR values in the same index position as the reference_tensor. For example, given a reference tensor with shape [B, 3, H, W] and setting bgr_tensor_shape_position == -3 indicates that the BGR position in this reference_tensor is at shape index -3, which is equivalent to index 1. Given these inputs, the resulting BGR tensor will have shape [1, 3, 1, 1], and the BGR values will be at shape index 1.
- Parameters:
- bgr_valUnion[float, Tuple[float, float, float], np.ndarray, torch.Tensor]
The BGR values to be converted into a tensor that will have a compatible shape with the input.
- reference_tensor: torch.Tensor
The tensor with the reference shape to convert the bgr_val.
- bgr_tensor_shape_positionint, default -3
Which position of the reference_tensor corresponds to the BGR values. Typical values are -1 (used in channels-last tensors) or -3 (…, CHW tensors)
- Returns:
- torch.Tensor
The bgr_val converted into a tensor with a shape compatible with reference_tensor.
- ptlflow.utils.utils.count_parameters(model: Module) int[source]
Return the number of trainable parameters of a model.
Taken from: https://discuss.pytorch.org/t/how-do-i-check-the-number-of-parameters-of-a-model/4325/7.
- Parameters:
- modeltorch.nn.Module
The model to count the parameters from.
- Returns:
- int
The number of trainable parameters of the given model.
- ptlflow.utils.utils.forward_interpolate_batch(prev_flow: Tensor) Tensor[source]
Apply RAFT’s forward_interpolate in a batch of torch.Tensors.
forward_interpolate in the warm start strategy where the previous flow estimation is forward projected and then used as initialization for the next estimation.
- Parameters:
- prev_flowtorch.Tensor
A 4D tensor [B, 2, H, W] containing a batch of previous flow predictions.
- Returns:
- torch.Tensor
The previous flow predictions after being forward interpolated.
- ptlflow.utils.utils.get_list_of_available_models_list() List[str][source]
Return a list of the names of the available models.
- Returns:
- list[str]
The list with the model names.
- ptlflow.utils.utils.make_divisible(v: int, div: int) int[source]
Decrease a number v until it is divisible by div.
- Parameters:
- vint
The number to be made divisible.
- divint
The divisor value.
- Returns:
- int
The new value of v which is divisible by div.
- ptlflow.utils.utils.release_gpu(tensors_dict: Dict[str, Any]) Dict[str, Any][source]
Detach and move to cpu the tensors from the input dict.
The non-tensor elements are kept intact.
- Parameters:
- tensors_dictDict[str, Any]
A dictionary containing the tensors to move.
- Returns:
- Dict[str, Any]
The same dictionary, but with the tensors moved to cpu.
- ptlflow.utils.utils.tensor_dict_to_numpy(tensor_dict: Dict[str, Tensor], padder: InputPadder | None = None) Dict[str, ndarray][source]
Convert all tensors into numpy format, changing the shape from CHW to HWC.
If “flows” is available, then a color representation “flows_viz” is added to the outputs.
- Parameters:
- tensor_dictdict[str, torch.Tensor]
A dictionary with the torch.Tensor inputs/outputs of the model.
- padder: InputPadder
Helper to unpad the images back to their original sizes.
- Returns:
- dict[str, np.ndarray]
The torch.Tensor entries from tensor_dict converted to numpy format.