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