|
|
from typing import List, Iterable
|
|
|
import torch
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
|
|
|
|
|
|
def pad_divide_by(in_img: torch.Tensor, d: int) -> (torch.Tensor, Iterable[int]):
|
|
|
h, w = in_img.shape[-2:]
|
|
|
|
|
|
if h % d > 0:
|
|
|
new_h = h + d - h % d
|
|
|
else:
|
|
|
new_h = h
|
|
|
if w % d > 0:
|
|
|
new_w = w + d - w % d
|
|
|
else:
|
|
|
new_w = w
|
|
|
lh, uh = int((new_h - h) / 2), int(new_h - h) - int((new_h - h) / 2)
|
|
|
lw, uw = int((new_w - w) / 2), int(new_w - w) - int((new_w - w) / 2)
|
|
|
pad_array = (int(lw), int(uw), int(lh), int(uh))
|
|
|
out = F.pad(in_img, pad_array)
|
|
|
return out, pad_array
|
|
|
|
|
|
|
|
|
def unpad(img: torch.Tensor, pad: Iterable[int]) -> torch.Tensor:
|
|
|
if len(img.shape) == 4:
|
|
|
if pad[2] + pad[3] > 0:
|
|
|
img = img[:, :, pad[2]:-pad[3], :]
|
|
|
if pad[0] + pad[1] > 0:
|
|
|
img = img[:, :, :, pad[0]:-pad[1]]
|
|
|
elif len(img.shape) == 3:
|
|
|
if pad[2] + pad[3] > 0:
|
|
|
img = img[:, pad[2]:-pad[3], :]
|
|
|
if pad[0] + pad[1] > 0:
|
|
|
img = img[:, :, pad[0]:-pad[1]]
|
|
|
elif len(img.shape) == 5:
|
|
|
if pad[2] + pad[3] > 0:
|
|
|
img = img[:, :, :, pad[2]:-pad[3], :]
|
|
|
if pad[0] + pad[1] > 0:
|
|
|
img = img[:, :, :, :, pad[0]:-pad[1]]
|
|
|
else:
|
|
|
raise NotImplementedError
|
|
|
return img
|
|
|
|
|
|
|
|
|
|
|
|
def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
|
|
|
with torch.amp.autocast("cuda"):
|
|
|
prob = prob.float()
|
|
|
new_prob = torch.cat([torch.prod(1 - prob, dim=dim, keepdim=True), prob],
|
|
|
dim).clamp(1e-7, 1 - 1e-7)
|
|
|
logits = torch.log((new_prob / (1 - new_prob)))
|
|
|
|
|
|
return logits
|
|
|
|
|
|
|
|
|
|
|
|
def cls_to_one_hot(cls_gt: torch.Tensor, num_objects: int) -> torch.Tensor:
|
|
|
|
|
|
B, _, H, W = cls_gt.shape
|
|
|
one_hot = torch.zeros(B, num_objects + 1, H, W, device=cls_gt.device).scatter_(1, cls_gt, 1)
|
|
|
return one_hot |
