ls_mlkit.util.mask.bio_masker module

class ls_mlkit.util.mask.bio_masker.BioBackboneFrameMasker(**kwargs: dict[Any, Any])

Bases: MaskerInterface

apply_inpainting_mask(x_0: T, x_t: T, inpainting_mask: Tensor) → Tensor

1 represents the region that can be seen

apply_mask(x: T, mask: Tensor) → Tensor

Parameters:

• x – T, (b,n,3,3) rotation matrix and (b,n,3) translation vector

• mask – b, n

check_mask_shape(x: T, mask: Tensor)

check whether the shape of mask is as expected

count_bright_area(mask: Tensor) → Tensor

Bright area can be seen Dark area cannot be seen

get_full_bright_mask(x: T) → Tensor

b, n, 3 -> b, n

class ls_mlkit.util.mask.bio_masker.BioCAOnlyMasker(ndim_mini_micro_shape: int = 1, **kwargs: dict[Any, Any])

Bases: MaskerInterface

apply_inpainting_mask(x_0: Tensor, x_t: Tensor, inpainting_mask: Tensor) → Tensor

1 represents the region that can be seen

apply_mask(x: Tensor, mask: Tensor) → Tensor

check_mask_shape(x: Tensor, mask: Tensor)

check whether the shape of mask is as expected

count_bright_area(mask: Tensor) → Tensor

Bright area can be seen Dark area cannot be seen

get_full_bright_mask(x: Tensor) → Tensor

b, n, 3 -> b, n

class ls_mlkit.util.mask.bio_masker.BioSO3Masker(**kwargs: dict[Any, Any])

Bases: MaskerInterface

apply_inpainting_mask(x_0: Tensor, x_t: Tensor, inpainting_mask: Tensor) → Tensor

1 represents the region that can be seen

apply_mask(x: Tensor, mask: Tensor) → Tensor

Parameters:

• x – (b,n,3,3) rotation matrix

• mask – b, n

check_mask_shape(x: Tensor, mask: Tensor)

check whether the shape of mask is as expected

count_bright_area(mask: Tensor) → Tensor

Bright area can be seen Dark area cannot be seen

get_full_bright_mask(x: Tensor) → Tensor

b, n, 3 -> b, n