-
Notifications
You must be signed in to change notification settings - Fork 83
/
post_prob.py
52 lines (45 loc) · 2.08 KB
/
post_prob.py
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
import torch
from torch.nn import Module
class Post_Prob(Module):
def __init__(self, sigma, c_size, stride, background_ratio, use_background, device):
super(Post_Prob, self).__init__()
assert c_size % stride == 0
self.sigma = sigma
self.bg_ratio = background_ratio
self.device = device
# coordinate is same to image space, set to constant since crop size is same
self.cood = torch.arange(0, c_size, step=stride,
dtype=torch.float32, device=device) + stride / 2
self.cood.unsqueeze_(0)
self.softmax = torch.nn.Softmax(dim=0)
self.use_bg = use_background
def forward(self, points, st_sizes):
num_points_per_image = [len(points_per_image) for points_per_image in points]
all_points = torch.cat(points, dim=0)
if len(all_points) > 0:
x = all_points[:, 0].unsqueeze_(1)
y = all_points[:, 1].unsqueeze_(1)
x_dis = -2 * torch.matmul(x, self.cood) + x * x + self.cood * self.cood
y_dis = -2 * torch.matmul(y, self.cood) + y * y + self.cood * self.cood
y_dis.unsqueeze_(2)
x_dis.unsqueeze_(1)
dis = y_dis + x_dis
dis = dis.view((dis.size(0), -1))
dis_list = torch.split(dis, num_points_per_image)
prob_list = []
for dis, st_size in zip(dis_list, st_sizes):
if len(dis) > 0:
if self.use_bg:
min_dis = torch.clamp(torch.min(dis, dim=0, keepdim=True)[0], min=0.0)
bg_dis = (st_size * self.bg_ratio) ** 2 / (min_dis + 1e-5)
dis = torch.cat([dis, bg_dis], 0) # concatenate background distance to the last
dis = -dis / (2.0 * self.sigma ** 2)
prob = self.softmax(dis)
else:
prob = None
prob_list.append(prob)
else:
prob_list = []
for _ in range(len(points)):
prob_list.append(None)
return prob_list