Advertisement
Guest User

Untitled

a guest
Jul 23rd, 2019
90
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.25 KB | None | 0 0
  1. def image_to_probs(image: np.ndarray, device=None):
  2.     """
  3.    Преобразует картинку в маски соответствующие классам
  4.    """
  5.     result = torch.zeros((len(channels),) + image.shape[-3:-1], device=device, dtype=torch.uint8) # матрица с нулями, например (4, 720, 480)
  6.     image = torch.tensor(np.moveaxis(image, -1, -3), device=device) > 128 # (3, 720, 1280) матрица типа bool
  7.     for channel, color in enumerate(channels):
  8.         color = torch.tensor(color, dtype=torch.uint8, device=device).unsqueeze(-1).unsqueeze(-1) # повернутая матрица каналов
  9.         mask = (image == color).all(dim=-3) # bool матрица, показывает соответствует ли пиксель классу
  10.         result[channel, :, :] = mask # соответствие канала, маске
  11.     result[0, result.any(dim=-3) == 0] = 1 # ищет пиксель, который не соответствует ни одному классу (неизветно как используется)
  12.     return result.to(torch.float32) # 8 слоев с заполнеными единицами пикселями сооцветсвующие данному классу
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement