23 lines
514 B
Python
23 lines
514 B
Python
import torch
|
|
|
|
|
|
def avg_color(img: torch.Tensor):
|
|
return img.mean(dim=(1, 2))
|
|
|
|
|
|
def mask(
|
|
foreground: torch.Tensor,
|
|
background: torch.Tensor,
|
|
mask_tensor: torch.Tensor,
|
|
threshold: float,
|
|
):
|
|
mask = mask_tensor > threshold
|
|
if foreground.dim() == 3:
|
|
mask = mask.unsqueeze(0)
|
|
mask = mask.expand(foreground.shape)
|
|
return torch.where(mask, foreground, background)
|
|
|
|
|
|
def add_matrix_vector(matrix: torch.Tensor, vector: torch.Tensor):
|
|
return matrix + vector.unsqueeze(1)
|