lib.torch_utils Module

Common multi-backend Torch utilities

class lib.torch_utils.ColorSpaceConvert(from_space: Literal['srgb', 'rgb', 'ycxcz', 'xyz'], to_space: Literal['lab', 'rgb', 'ycxcz', 'xyz'])

Transforms inputs between different color spaces on the GPU. Images expected in (N,C,H,W) order

Notes

The following color space transformations are implemented:

  • rgb to lab

  • rgb to xyz

  • srgb to _rgb

  • srgb to ycxcz

  • xyz to ycxcz

  • xyz to lab

  • xyz to rgb

  • ycxcz to rgb

  • ycxcz to xyz

Parameters:

  • from_space (T.Literal['srgb', 'rgb', 'ycxcz', 'xyz']) – One of “srgb”, “rgb”, “ycxcz”, “xyz”

  • to_space (T.Literal['lab', 'rgb', 'ycxcz', 'xyz']) – One of “lab”, “rgb”, “ycxcz”, “xyz”

Raises:

ValueError – If the requested color space conversion is not defined

forward(image: Tensor) → Tensor

Call the color-space conversion function.

Parameters:

image (Tensor) – The image tensor in the color-space defined by from_space

Return type:

The image tensor in the color-space defined by to_space

lib.torch_utils.get_device(cpu: bool = False) → device

Get the correctly configured device for running Torch

Parameters:

cpu (bool) – True to force running on the CPU.

Return type:

The device that torch should use

Functions

get_device([cpu])

Get the correctly configured device for running Torch

Classes

ColorSpaceConvert(from_space, to_space)

Transforms inputs between different color spaces on the GPU.