Spaces:
Running
on
A10G
Running
on
A10G
File size: 1,942 Bytes
320e465 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
from typing import List, Iterable
import torch
import torch.nn.functional as F
# STM
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
# @torch.jit.script
def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
with torch.cuda.amp.autocast(enabled=False):
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
# @torch.jit.script
def cls_to_one_hot(cls_gt: torch.Tensor, num_objects: int) -> torch.Tensor:
# cls_gt: B*1*H*W
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 |