36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import torch
|
|
|
|
|
|
def get_mask(
|
|
foreground_mask: torch.Tensor | None = None,
|
|
foreground_MASK: torch.Tensor | None = None,
|
|
):
|
|
if foreground_mask is None and foreground_MASK is None:
|
|
raise ValueError("Please provide one mask image")
|
|
|
|
if foreground_MASK is not None:
|
|
mask = foreground_MASK.squeeze()
|
|
else:
|
|
mask = img_to_mask(foreground_mask.permute(0, 3, 1, 2)).squeeze()
|
|
return mask
|
|
|
|
|
|
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
|
|
|
|
|
|
def get_screen(batch_size: int, height: int, width: int, r: int, g: int, b: int):
|
|
# Normalize RGB values to the range 0.0-1.0
|
|
r_normalized = float(r) / 255.0
|
|
g_normalized = float(g) / 255.0
|
|
b_normalized = float(b) / 255.0
|
|
|
|
rgb_image = torch.zeros((batch_size, height, width, 3), dtype=torch.float32)
|
|
rgb_image[:, :, :, 0] = r_normalized # Red channel (index 0)
|
|
rgb_image[:, :, :, 1] = g_normalized # Green channel (index 1)
|
|
rgb_image[:, :, :, 2] = b_normalized # Blue channel (index 2)
|
|
return rgb_image
|