def test_multioutput2(self):
a, b = jt.index([3, 3])
assert (a.data == [[0, 0, 0], [1, 1, 1], [2, 2, 2]]).all()
assert (b.data == [[0, 1, 2], [0, 1, 2], [0, 1, 2]]).all(), b.data
a, b = jt.index([3, 3])
c = a + b
assert (c.data == [[0, 1, 2], [1, 2, 3], [2, 3, 4]]).all(), c.data
def test(self):
assert (jt.index([2, 2], 0).data == [[0, 0], [1, 1]]).all()
assert (jt.index([2, 2], 1).data == [[0, 1], [0, 1]]).all()
a = jt.index([2, 2], 0)
b = jt.index([2, 2], 1)
c = a + b
assert (c.data == [[0, 1], [1, 2]]).all(), c.data
def grid_sampler_2d(X,grid,mode,padding_mode,align_corners):
N = X.shape[0]
C = X.shape[1]
inp_H = X.shape[2]
inp_W = X.shape[3]
H = grid.shape[1]
W = grid.shape[2]
x = grid[:,:,:,0]
y = grid[:,:,:,1]
shape = [N,C,H,W]
cid = jt.index(shape, dim=1)
nid = jt.index(shape, dim=0)
x = grid_sampler_compute_source_index(x,inp_W,padding_mode,align_corners)
y = grid_sampler_compute_source_index(y,inp_H,padding_mode,align_corners)
xid = x.reindex(shape,['i0','i2','i3'])
yid = y.reindex(shape,['i0','i2','i3'])
if mode=='nearest':
return X.reindex([nid,cid,yid.round(),xid.round()])
elif mode=='bilinear':
#xid,yid = (xid+0.00001),(yid+0.00001)
fx,fy = (xid).floor(),(yid).floor()
cx,cy = fx+1,fy+1
dx,dy = xid-fx,yid-fy
dnx,dny = cx-xid,cy-yid
a = X.reindex([nid,cid,fy,fx],overflow_value=0.0)
b = X.reindex([nid,cid,cy,fx],overflow_value=0.0)
c = X.reindex([nid,cid,fy,cx],overflow_value=0.0)
d = X.reindex([nid,cid,cy,cx],overflow_value=0.0)
o = a*dnx*dny+b*dnx*dy+c*dx*dny+d*dx*dy
return o
def __call__(self, proposals, mask_logits, targets):
"""
Arguments:
proposals (list[BoxList])
mask_logits (Tensor)
targets (list[BoxList])
Return:
mask_loss (Tensor): scalar tensor containing the loss
"""
labels, mask_targets, mask_ratios = self.prepare_targets(
proposals, targets)
labels = cat(labels, dim=0)
mask_targets = cat(mask_targets, dim=0)
positive_inds = jt.nonzero(labels > 0).squeeze(1)
labels_pos = labels[positive_inds]
# accept empty tensors, so handle it separately
if mask_targets.numel() == 0:
if not self.maskiou_on:
return mask_logits.sum() * 0
else:
selected_index = jt.arange(mask_logits.shape[0])
selected_mask = mask_logits[selected_index, labels]
mask_num, mask_h, mask_w = selected_mask.shape
selected_mask = selected_mask.reshape(mask_num, 1, mask_h,
mask_w)
return mask_logits.sum() * 0, selected_mask, labels, None
if self.maskiou_on:
mask_ratios = cat(mask_ratios, dim=0)
value_eps = 1e-10 * jt.ones((mask_targets.shape[0], ))
mask_ratios = jt.maximum(mask_ratios, value_eps)
pred_masks = mask_logits[positive_inds, labels_pos]
pred_masks[:] = pred_masks > 0
mask_targets_full_area = mask_targets.sum(
dims=[1, 2]) / mask_ratios
mask_ovr = pred_masks * mask_targets
mask_ovr_area = mask_ovr.sum(dims=[1, 2])
mask_union_area = pred_masks.sum(
dims=[1, 2]) + mask_targets_full_area - mask_ovr_area
value_1 = jt.ones((pred_masks.shape[0], ))
value_0 = jt.zeros((pred_masks.shape[0], ))
mask_union_area = jt.maximum(mask_union_area, value_1)
mask_ovr_area = jt.maximum(mask_ovr_area, value_0)
maskiou_targets = mask_ovr_area / mask_union_area
binary_cross_entropy_with_logits = nn.BCEWithLogitsLoss()
mask_loss = binary_cross_entropy_with_logits(
mask_logits[positive_inds, labels_pos], mask_targets)
if not self.maskiou_on:
return mask_loss
else:
selected_index = jt.index((mask_logits.shape[0], ), dim=0)
selected_mask = mask_logits[selected_index, labels]
mask_num, mask_h, mask_w = selected_mask.shape
selected_mask = selected_mask.reshape(mask_num, 1, mask_h, mask_w)
selected_mask = selected_mask.sigmoid()
return mask_loss, selected_mask, labels, maskiou_targets
def check_cub_argsort(shape, dim, descending = False):
with jt.log_capture_scope(
log_silent=1,
log_v=0, log_vprefix="op.cc=100"
) as raw_log:
x = jt.random(shape)
y, y_key = jt.argsort(x, dim=dim, descending=descending)
v = []
for i in range(len(shape)):
if (i == dim):
v.append(y)
else:
v.append(jt.index(shape, dim=i))
yk = jt.reindex(x, v)
yk_ = yk.data
y_key_ = y_key.data
logs = find_log_with_re(raw_log, "(Jit op key (not )?found: " + "cub_argsort" + ".*)")
assert len(logs)==1
x__ = x.data
if descending:
x__ = -x__
yk__ = np.sort(x__, axis=dim)
if descending:
yk__ = -yk__
assert np.allclose(y_key_, yk__)
assert np.allclose(yk_, yk__)
def grid_sampler_3d(X, grid, mode, padding_mode, align_corners):
N = X.shape[0]
C = X.shape[1]
inp_D = X.shape[2]
inp_H = X.shape[3]
inp_W = X.shape[4]
D = grid.shape[1]
H = grid.shape[2]
W = grid.shape[3]
x = grid[:, :, :, :, 0]
y = grid[:, :, :, :, 1]
z = grid[:, :, :, :, 2]
shape = [N, C, D, H, W]
cid = jt.index(shape, dim=1)
nid = jt.index(shape, dim=0)
x = grid_sampler_compute_source_index(x, inp_W, padding_mode,
align_corners)
y = grid_sampler_compute_source_index(y, inp_H, padding_mode,
align_corners)
z = grid_sampler_compute_source_index(z, inp_D, padding_mode,
align_corners)
xid = x.reindex(shape, ['i0', 'i2', 'i3', 'i4'])
yid = y.reindex(shape, ['i0', 'i2', 'i3', 'i4'])
zid = z.reindex(shape, ['i0', 'i2', 'i3', 'i4'])
if mode == 'nearest':
return X.reindex([nid, cid, zid.round(), yid.round(), xid.round()])
elif mode == 'bilinear':
fx, fy, fz = xid.floor(), yid.floor(), zid.floor()
cx, cy, cz = fx + 1, fy + 1, fz + 1
dx, dy, dz = xid - fx, yid - fy, zid - fz
dnx, dny, dnz = cx - xid, cy - yid, cz - zid
a = X.reindex([nid, cid, fz, fy, fx])
b = X.reindex([nid, cid, cz, fy, fx])
c = X.reindex([nid, cid, fz, cy, fx])
d = X.reindex([nid, cid, fz, fy, cx])
e = X.reindex([nid, cid, fz, cy, cx])
f = X.reindex([nid, cid, cz, fy, cx])
g = X.reindex([nid, cid, cz, cy, fx])
h = X.reindex([nid, cid, cz, cy, cx])
o = a * dnx * dny * dnz + b * dnx * dny * dz + c * dnx * dy * dnz + d * dx * dny * dnz + e * dx * dy * dnz + f * dx * dny * dz + g * dnx * dy * dz + h * dx * dy * dz
return o
def upsample(img, size, mode="nearest", align_corners=False):
n, c, h, w = img.shape
H, W = size
nid, cid, hid, wid = jt.index((n, c, H, W))
if align_corners:
x = hid * ((h - 1) / max(1, H - 1))
y = wid * ((w - 1) / max(1, W - 1))
else:
x = hid * (h / H)
y = wid * (w / W)
return _interpolate(img, x, y, (nid, cid), mode)
def arange(start=0, end=None, step=1,dtype=None):
if end is None:
end,start = start,0
l = round((end-start)//step)+1
if (l-1)*step+start>=end:
l-=1
x = jt.index((l,),0)
x = x*step+start
if dtype is not None:
x= x.cast(dtype)
return x
def resize_and_crop(x, bbox, interpolation="nearest", out_size=[224, 224]):
N, k = bbox.shape
H, W, C = x.shape
assert k == 4
shape = [N, out_size[0], out_size[1], C]
# shape = [N,H,W]
# fx x cx
# +------------>
# fy | a dx | b
# | dy
# y | - o -
# |
# cy | c | d
# v
img = x
bb = [bbox.reindex(shape, ["i0", str(i)]) for i in range(4)]
hid = jt.index(shape, dim=1)
wid = jt.index(shape, dim=2)
cid = jt.index(shape, dim=3)
one = jt.array(1.0).broadcast(shape)
x = bb[0] * (H - 1.0) + hid * ((H - 1) * 1.0 /
(shape[1] - 1)) * (bb[2] - bb[0])
y = bb[1] * (W - 1.0) + wid * ((W - 1) * 1.0 /
(shape[2] - 1)) * (bb[3] - bb[1])
if interpolation == "nearest":
return img.reindex([x.round(), y.round(), cid])
if interpolation == "bilinear":
fx, fy = x.floor(), y.floor()
cx, cy = fx + one, fy + one
dx, dy = x - fx, y - fy
a = img.reindex_var([fx, fy, cid])
b = img.reindex_var([cx, fy, cid])
c = img.reindex_var([fx, cy, cid])
d = img.reindex_var([cx, cy, cid])
dnx, dny = one - dx, one - dy
ab = dx * b + dnx * a
cd = dx * d + dnx * c
o = ab * dny + cd * dy
return o
raise (f"Not support {interpolation}")
def resize(img, size, mode="nearest", align_corners=False):
n, c, h, w = img.shape
H, W = size
nid, cid, hid, wid = jt.index((n, c, H, W))
if align_corners:
x = hid * ((h - 1) / max(1, H - 1))
y = wid * ((w - 1) / max(1, W - 1))
else:
x = hid * (h / H) + (h / H * 0.5 - 0.5)
if H > h: x = x.clamp(0, h - 1)
y = wid * (w / W) + (w / W * 0.5 - 0.5)
if W > w: y = y.clamp(0, w - 1)
return _interpolate(img, x, y, (nid, cid), mode)
def unique(x):
r'''
Returns the unique elements of the input tensor.
Args:
x– the input tensor.
'''
x = x.reshape(-1)
_, x = jt.argsort(x)
index, = jt.index((x.shape[0], ))
y = x[1:][x[index[1:]] != x[index[:-1]]]
x = jt.contrib.concat([x[:1], y], dim=0)
return x
def gather(x,dim,index):
if dim<0:
dim+=index.ndim
x_shape = list(x.shape )
i_shape = list(index.shape)
assert i_shape[dim]>0
assert x.ndim == index.ndim
i_shape[dim]=x_shape[dim]
assert i_shape == x_shape
ins = []
for i in range(index.ndim):
ins.append(jt.index(index.shape,dim=i))
ins[dim]=index
return x.reindex(ins)
def resize_and_crop(x, bbox, interpolation="nearest"):
N, k = bbox.shape
H, W = x.shape
assert k==4
shape = [N,H,W]
# fx x cx
# +------------>
# fy | a dx | b
# | dy
# y | - o -
# |
# cy | c | d
# v
img = x
bb = [ bbox.reindex(shape, ["i0", str(i)]) for i in range(4) ]
hid = jt.index(shape, 1)
wid = jt.index(shape, 2)
one = jt.float(1).broadcast(shape)
x = bb[0]*jt.float(H-1)+hid*(bb[2]-bb[0])
y = bb[1]*jt.float(W-1)+wid*(bb[3]-bb[1])
if interpolation=="nearest":
return img.reindex_var([x.round(), y.round()])
if interpolation=="bilinear":
fx, fy = x.floor(), y.floor()
cx, cy = fx+one, fy+one
dx, dy = x-fx, y-fy
a = img.reindex_var([fx, fy])
b = img.reindex_var([cx, fy])
c = img.reindex_var([fx, cy])
d = img.reindex_var([cx, cy])
dnx, dny = one-dx, one-dy
ab = dx*b + dnx*a
cd = dx*d + dnx*c
o = ab*dny + cd*dy
return o
raise(f"Not support {interpolation}")
def grid_sample_v0(input, grid, mode='bilinear', padding_mode='zeros'):
r'''
Given an input and a flow-field grid, computes the output using input values and pixel locations from grid.
grid specifies the sampling pixel locations normalized by the input spatial dimensions. Therefore, it should have most values in the range of [-1, 1]. For example, values x = -1, y = -1 is the left-top pixel of input, and values x = 1, y = 1 is the right-bottom pixel of input.
Args:
[in] input (var): the source input var, whose shape is (N, C, Hi, Wi)
[in] grid (var): the pixel locations, whose shape is (N, Ho, Wo, 2)
[in] mode (string): the interpolate way, default: bilinear.
[in] padding_mode (string): the padding way, default: zeros.
[out] output (var): the output var, whose shape is (N, C, Ho, Wo)
Example:
>>> x = jt.array([[[[1,2],[3,4]]]])
>>> print(x)
[[[[1 2]
[3 4]]]]
>>> grid = jt.array([[[[0.5, 0.5]]]])
>>> print(x.shape, grid.shape)
[1,1,2,2,], [1,1,2,2,]
>>> nn.grid_sample(x, grid)
[[[[3.25]]]]
'''
assert padding_mode == 'zeros'
Ni, Ci, Hi, Wi = input.shape
No, Ho, Wo, D = grid.shape
assert D == 2
assert Ni == No
assert len(input.shape) == 4 and len(grid.shape)
nid, cid, hid, wid = jt.index((Ni, Ci, Ho, Wo))
x = ((grid[:, :, :, 1].unsqueeze(1).repeat([1, Ci, 1, 1]) + 1) / 2) * (Hi -
1)
y = ((grid[:, :, :, 0].unsqueeze(1).repeat([1, Ci, 1, 1]) + 1) / 2) * (Wi -
1)
return _interpolate(input, x, y, (nid, cid), mode)
def check_argsort(shape, dim, descending = False):
x = jt.random(shape)
y, y_key = jt.argsort(x, dim=dim, descending=descending)
v = []
for i in range(len(shape)):
if (i == dim):
v.append(y)
else:
v.append(jt.index(shape, dim=i))
yk = jt.reindex(x, v)
yk_ = yk.data
y_key_ = y_key.data
x__ = x.data
if descending:
x__ = -x__
yk__ = np.sort(x__, axis=dim)
if descending:
yk__ = -yk__
assert np.allclose(y_key_, yk__)
assert np.allclose(yk_, yk__)
def execute(self, boxes, pred_maskiou, labels):
num_masks = pred_maskiou.shape[0]
index = jt.index((num_masks, ), 0)
maskious = pred_maskiou[index, labels]
# maskious = [maskious]
# split `maskiou` accroding to `boxes`
boxes_per_image = [len(box) for box in boxes]
maskious = maskious.split(boxes_per_image, dim=0)
# results = []
for maskiou, bbox in zip(maskious, boxes):
# bbox = BoxList(bbox.bbox, bbox.size, mode="xyxy")
# for field in bbox.fields():
# bbox.add_field(field, bbox.get_field(field))
bbox_scores = bbox.get_field("scores")
mask_scores = bbox_scores * maskiou
bbox.add_field("mask_scores", mask_scores)
# results.append(bbox)
# return results
return boxes
def resize(x, size, mode="nearest"):
img = x
n, c, h, w = x.shape
H, W = size
new_size = [n, c, H, W]
nid, cid, hid, wid = jt.index(new_size)
x = hid * ((h - 1) / (H - 1))
y = wid * ((w - 1) / (W - 1))
if mode == "nearest":
return img.reindex([nid, cid, x.floor(), y.floor()])
if mode == "bilinear":
fx, fy = x.floor(), y.floor()
cx, cy = fx + 1, fy + 1
dx, dy = x - fx, y - fy
a = img.reindex_var([nid, cid, fx, fy])
b = img.reindex_var([nid, cid, cx, fy])
c = img.reindex_var([nid, cid, fx, cy])
d = img.reindex_var([nid, cid, cx, cy])
dnx, dny = 1 - dx, 1 - dy
ab = dx * b + dnx * a
cd = dx * d + dnx * c
o = ab * dny + cd * dy
return o
raise (f"Not support {interpolation}")
def build_targets(p, targets, model):
# Build targets for compute_loss(), input targets(image,class,x,y,w,h)
det = model.model[-1] # Detect() module
na, nt = det.na, targets.shape[0] # number of anchors, targets
tcls, tbox, indices, anch = [], [], [], []
gain = jt.ones((7, )) # normalized to gridspace gain
ai = jt.index(
(na, ),
dim=0).float().view(na, 1).repeat(1,
nt) # same as .repeat_interleave(nt)
targets = jt.contrib.concat((targets.repeat(na, 1, 1), ai[:, :, None]),
2) # append anchor indices
g = 0.5 # bias
off = jt.array(
[
[0, 0],
# [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
], ).float() * g # offsets
for i in range(det.nl):
anchors = det.anchors[i]
gain[2:6] = jt.array(
[p[i].shape[3], p[i].shape[2], p[i].shape[3],
p[i].shape[2]]) # xyxy gain
# Match targets to anchors
t = targets * gain
if nt:
# Matches
r = t[:, :, 4:6] / anchors[:, None] # wh ratio
j = jt.maximum(r, 1. / r).max(2) < model.hyp['anchor_t'] # compare
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
t = t[j] # filter
# Offsets
gxy = t[:, 2:4] # grid xy
gxi = gain[jt.array([2, 3])] - gxy # inverse
# j, k = jt.logical_and((gxy % 1. < g), (gxy > 1.)).int().transpose(1,0).bool()
# l, m = jt.logical_and((gxi % 1. < g),(gxi > 1.)).int().transpose(1,0).bool()
jk = jt.logical_and((gxy % 1. < g), (gxy > 1.))
lm = jt.logical_and((gxi % 1. < g), (gxi > 1.))
j, k = jk[:, 0], jk[:, 1]
l, m = lm[:, 0], lm[:, 1]
j = jt.stack((jt.ones_like(j), ))
t = t.repeat((off.shape[0], 1, 1))[j]
offsets = (jt.zeros_like(gxy)[None] + off[:, None])[j]
else:
t = targets[0]
offsets = 0
# Define
b = t[:, 0].int32()
c = t[:, 1].int32() # image, class
gxy = t[:, 2:4] # grid xy
gwh = t[:, 4:6] # grid wh
gij = (gxy - offsets).int32()
gi, gj = gij[:, 0], gij[:, 1] # grid xy indices
# Append
a = t[:, 6].int32() # anchor indices
indices.append((b, a, gj.clamp(0, gain[3] - 1),
gi.clamp(0,
gain[2] - 1))) # image, anchor, grid indices
tbox.append(jt.contrib.concat((gxy - gij, gwh), 1)) # box
anch.append(anchors[a]) # anchors
tcls.append(c) # class
return tcls, tbox, indices, anch
def test_multioutput3(self):
a, b = jt.index([3, 3])
del a
assert (b.data == [[0, 1, 2], [0, 1, 2], [0, 1, 2]]).all(), b.data
def test_multioutput(self):
a, b = jt.index([2, 2])
jt.sync([a, b])
assert (a.data == [[0, 0], [1, 1]]).all()
assert (b.data == [[0, 1], [0, 1]]).all(), b.data
def linspace(start, end, steps):
res = jt.index((steps, ))[0]
res = res * (end - start) / float(steps - 1) + start
return res
