This repository has been archived by the owner on Jul 12, 2020. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 10
/
BoxIoU.lua
109 lines (83 loc) · 2.89 KB
/
BoxIoU.lua
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
require 'nn'
local box_utils = require 'box_utils'
local layer, parent = torch.class('nn.BoxIoU', 'nn.Module')
function layer:__init()
parent.__init(self)
self.area1 = torch.CudaTensor()
self.area2 = torch.CudaTensor()
self.overlap = torch.CudaTensor()
self.output = torch.CudaTensor()
self.savedIoUs = torch.CudaTensor()
end
-- Convert from (xc, yc, w, h) to (x0, y0, x1, y1)
--[[
local function convert_boxes(boxes)
local ret = boxes.new(#boxes)
local xc = boxes:select(3, 1)
local yc = boxes:select(3, 2)
local w = boxes:select(3, 3)
local h = boxes:select(3, 4)
local x0 = ret:select(3, 1)
local x1 = ret:select(3, 3)
local y0 = ret:select(3, 2)
local y1 = ret:select(3, 4)
x0:div(w, 2.0):mul(-1):add(xc)
x1:div(w, 2.0):add(xc)
y0:div(h, 2.0):mul(-1):add(yc)
y1:div(h, 2.0):add(yc)
return ret
end
--]]
function layer:updateOutput(input)
local box1 = input[1]
local box2 = input[2]
--
local N, B1, B2 = box1:size(1), box1:size(2), box2:size(2)
self.area1:cmul(box1[{{}, {}, 3}], box1[{{}, {}, 4}])
self.area2:cmul(box2[{{}, {}, 3}], box2[{{}, {}, 4}])
local area1_expand = self.area1:view(N, B1, 1):expand(N, B1, B2):cuda()
local area2_expand = self.area2:view(N, 1, B2):expand(N, B1, B2):cuda()
local convert_boxes = box_utils.xcycwh_to_x1y1x2y2
local box1_lohi = convert_boxes(box1) -- N x B1 x 4
local box2_lohi = convert_boxes(box2) -- N x B2 x 4
local box1_lohi_expand = box1_lohi:view(N, B1, 1, 4):expand(N, B1, B2, 4)
local box2_lohi_expand = box2_lohi:view(N, 1, B2, 4):expand(N, B1, B2, 4)
local x0 = torch.cmax(box1_lohi_expand:select(4, 1),
box2_lohi_expand:select(4, 1))
local y0 = torch.cmax(box1_lohi_expand:select(4, 2),
box2_lohi_expand:select(4, 2))
local x1 = torch.cmin(box1_lohi_expand:select(4, 3),
box2_lohi_expand:select(4, 3))
local y1 = torch.cmin(box1_lohi_expand:select(4, 4),
box2_lohi_expand:select(4, 4))
local w = (x1 - x0):cmax(0)
local h = (y1 - y0):cmax(0)
local intersection = torch.cmul(w, h)
self.output:add(area1_expand, -1, intersection)
self.output:add(area2_expand):pow(-1)
self.output:cmul(intersection)
local savedIoUs = self.savedIoUs;
local sIou = nil;
if self.output:size(1)==1 then
sIou = self.output:reshape(self.output:size(1));
else
sIou = torch.squeeze(self.output)
end
-- require('mobdebug').start(nill,8222)
if savedIoUs:nDimension() == 0 then
self.savedIoUs:mul(sIou,1)
else
self.savedIoUs = savedIoUs:cat(sIou)
end
return self.output
end
function layer:updateGradInput(input, gradOutput)
error('Not implemented')
end
function layer:cleanIoUs()
self.savedIoUs = torch.CudaTensor();
end
function layer:avgIoU()
-- require('mobdebug').start(nill,8222);
return torch.sum(self.savedIoUs)/self.savedIoUs:nElement();
end