18 lines
516 B
Python
18 lines
516 B
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def ensure_size(tensor: torch.Tensor, size: tuple[int, int] = (128, 128)):
|
|
if tensor.shape[-2:] == size:
|
|
return tensor
|
|
else:
|
|
return F.interpolate(tensor, size=size, mode="bilinear", align_corners=False)
|
|
|
|
|
|
def img_to_mask(tensor: torch.Tensor):
|
|
weights = torch.tensor([0.2989, 0.5870, 0.1140], device=tensor.device)
|
|
weights = weights.view(1, 3, 1, 1)
|
|
grayscale = torch.sum(tensor * weights, dim=1, keepdim=True)
|
|
return grayscale
|
|
|