def __getitem__(self, index): A_path = self.files_A[index % len(self.files_A)] image_A = Image.open(A_path) # Convert grayscale images to rgb if image_A.mode != "RGB": image_A = to_rgb(image_A) item_A = self.transform(image_A) item_A = jt.array(item_A) item_A_l = {} regions = ['eyel','eyer','nose','mouth'] basen = os.path.basename(A_path)[:-4] lm_path = os.path.join(self.lmdir, basen+'.txt') feats = getfeats(lm_path) mouth_x = int((feats[3,0]+feats[4,0])/2.0) mouth_y = int((feats[3,1]+feats[4,1])/2.0) ratio = self.load_h // 256 rhs = np.array([EYE_H,EYE_H,NOSE_H,MOUTH_H]) * ratio rws = np.array([EYE_W,EYE_W,NOSE_W,MOUTH_W]) * ratio center = np.array([[feats[0,0],feats[0,1]-4*ratio],[feats[1,0],feats[1,1]-4*ratio],[feats[2,0],feats[2,1]-rhs[2]//2+16*ratio],[mouth_x,mouth_y]]) soft_border_mask4 = [] for i in range(4): xb = [np.zeros(rhs[i]),np.ones(rhs[i])*(rws[i]-1)] yb = [np.zeros(rws[i]),np.ones(rws[i])*(rhs[i]-1)] soft_border_mask = getSoft([rhs[i],rws[i]],xb,yb) soft_border_mask = jt.array(soft_border_mask).unsqueeze(0).float() soft_border_mask4.append(soft_border_mask) for i in range(4): item_A_l[regions[i]+'_A'] = item_A[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] * soft_border_mask4[i].repeat(3,1,1) cmasks = [] for i in range(4): cmaskpath = os.path.join(self.cmaskdir.format(regions[i]),basen+'.png') im_cmask = Image.open(cmaskpath) cmask0 = self.transform_mask(im_cmask) cmask0 = jt.array(cmask0) cmask0 = (cmask0 >= 0.5).float() cmask = cmask0.clone() cmask = cmask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] cmasks.append(cmask) mask = jt.ones([1,item_A.shape[1],item_A.shape[2]]) # mask out eyes, nose, mouth for i in range(4): mask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] = 0 imgsize = self.load_h maskn = mask[0].numpy() masks = [np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize])] masks[0][1:] = maskn[:-1] masks[1][:-1] = maskn[1:] masks[2][:,1:] = maskn[:,:-1] masks[3][:,:-1] = maskn[:,1:] masks2 = [maskn-e for e in masks] bound = np.minimum.reduce(masks2) bound = -bound xb = [] yb = [] for i in range(4): xbi = [center[i,0]-rws[i]/2, center[i,0]+rws[i]/2-1] ybi = [center[i,1]-rhs[i]/2, center[i,1]+rhs[i]/2-1] for j in range(2): maskx = bound[:,int(xbi[j])] masky = bound[int(ybi[j]),:] tmp_a = maskx * xbi[j] tmp_b = 1-maskx xb += [tmp_b*10000 + tmp_a] tmp_a = masky * ybi[j] tmp_b = 1-masky yb += [tmp_b*10000 + tmp_a] soft = 1-getSoft([imgsize,imgsize],xb,yb) soft = jt.array(soft).unsqueeze(0).float() mask = (jt.ones(mask.shape)-mask)*soft + mask bgpath = os.path.join(self.maskdir, basen+'.png') im_bg = Image.open(bgpath) mask2 = self.transform_mask(im_bg) # mask out background mask2 = jt.array(mask2) mask2 = (mask2 >= 0.5).float() # foreground: 1, background: 0 item_A_l['hair_A'] = (item_A/2+0.5) * mask.repeat(3,1,1) * mask2.repeat(3,1,1) * 2 - 1 item_A_l['bg_A'] = (item_A/2+0.5) * (jt.ones(mask2.shape)-mask2).repeat(3,1,1) * 2 - 1 return item_A, item_A_l['eyel_A'], item_A_l['eyer_A'], item_A_l['nose_A'], item_A_l['mouth_A'], item_A_l['hair_A'], item_A_l['bg_A'], mask, mask2, center, cmasks[0], cmasks[1], cmasks[2], cmasks[3]
def test_setitem_inplace_case2(self): # test un-continuous first dim a = jt.zeros((3, )) a[0::2] = jt.ones((2, )) assert a.data[2] == 1
def main(): model = PSPNet() x = jt.ones([2, 3, 513, 513]) y = model(x) print(y.shape) _ = y.data
np.array([num for _ in range(n_row) for num in range(n_row)])).float32().stop_grad() gen_imgs = generator(z, labels) save_image(gen_imgs.numpy(), "images/%d.png" % batches_done, nrow=n_row) # ---------- # Training # ---------- for epoch in range(opt.n_epochs): for i, (imgs, labels) in enumerate(dataloader): batch_size = imgs.shape[0] # Adversarial ground truths valid = jt.ones([batch_size, 1]).float32().stop_grad() fake = jt.zeros([batch_size, 1]).float32().stop_grad() # Configure input real_imgs = jt.array(imgs) labels = jt.array(labels) # ----------------- # Train Generator # ----------------- # Sample noise and labels as generator input z = jt.array(np.random.normal(0, 1, (batch_size, opt.latent_dim))).float32() gen_labels = jt.array(np.random.randint(0, opt.n_classes, batch_size)).float32()
def read_planetoid_data(folder, prefix): names = ['x', 'tx', 'allx', 'y', 'ty', 'ally', 'graph', 'test.index'] items = [read_file(folder, prefix, name) for name in names] x, tx, allx, y, ty, ally, graph, test_index = items train_index = jt.arange(y.size(0), dtype=Var.int32) val_index = jt.arange(y.size(0), y.size(0) + 500, dtype=Var.int32) sorted_test_index = test_index.argsort()[1] if prefix.lower() == 'citeseer': # There are some isolated nodes in the Citeseer graph, resulting in # none consecutive test indices. We need to identify them and add them # as zero vectors to `tx` and `ty`. len_test_indices = (test_index.max() - test_index.min()).item() + 1 tx_ext = jt.zeros((len_test_indices, tx.size(1))) tx_ext[sorted_test_index - test_index.min(), :] = tx ty_ext = jt.zeros((len_test_indices, ty.size(1))) ty_ext[sorted_test_index - test_index.min(), :] = ty tx, ty = tx_ext, ty_ext if prefix.lower() == 'nell.0.001': tx_ext = jt.zeros((len(graph) - allx.size(0), x.size(1))) tx_ext[sorted_test_index - allx.size(0)] = tx ty_ext = jt.zeros((len(graph) - ally.size(0), y.size(1))) ty_ext[sorted_test_index - ally.size(0)] = ty tx, ty = tx_ext, ty_ext x = jt.concat([allx, tx], dim=0) x[test_index] = x[sorted_test_index] # Creating feature vectors for relations. row, col, value = SparseTensor.from_dense(x).coo() rows, cols, values = [row], [col], [value] mask1 = index_to_mask(test_index, size=len(graph)) mask2 = index_to_mask(jt.arange(allx.size(0), len(graph)), size=len(graph)) mask = jt.logical_or(jt.logical_not(mask1), jt.logical_not(mask2)) isolated_index = mask.nonzero(as_tuple=False).view(-1)[allx.size(0):] rows += [isolated_index] cols += [jt.arange(isolated_index.size(0)) + x.size(1)] values += [jt.ones((isolated_index.size(0)))] x = SparseTensor(row=jt.concat(rows), col=jt.concat(cols), value=jt.concat(values)) else: x = jt.concat([allx, tx], dim=0) x[test_index] = x[sorted_test_index] y = jt.concat([ally, ty], dim=0).argmax(dim=1)[0] y[test_index] = y[sorted_test_index] train_mask = index_to_mask(train_index, size=y.size(0)) val_mask = index_to_mask(val_index, size=y.size(0)) test_mask = index_to_mask(test_index, size=y.size(0)) edge_index = edge_index_from_dict(graph, num_nodes=y.size(0)) data = Data(x=x, edge_index=edge_index, y=y) data.train_mask = train_mask data.val_mask = val_mask data.test_mask = test_mask return data
def sign(x): one = jt.ones(x.shape) x = jt.ternary(x > 0, one, x) return jt.ternary(x < 0, -one, x)
def addone_with_mask(A, mask): return ((A/2+0.5)*mask + (jt.ones(mask.shape)-mask))*2-1
def test_view(self): a = jt.ones([2,3,4]) assert a.view(2,-1).shape == [2,12]
def test_permute(self): a = jt.ones([2, 3, 4]) assert a.permute().shape == [4, 3, 2] assert a.permute(0, 2, 1).shape == [2, 4, 3]
def inverse_mask(mask): return jt.ones(mask.shape)-mask
def execute(self, x, img_size, scale=1.0): """Forward Region Proposal Network. Here are notations. * :math:`N` is batch size. * :math:`C` channel size of the input. * :math:`H` and :math:`W` are height and witdh of the input feature. * :math:`A` is number of anchors assigned to each pixel. Args: x : The Features extracted from images. Its shape is :math:`(N, C, H, W)`. img_size (tuple of ints): A tuple :obj:`height, width`, which contains image size after scaling. Returns: This is a tuple of five following values. * **rpn_locs**: Predicted bounding box offsets and scales for \ anchors. Its shape is :math:`(N, H W A, 4)`. * **rpn_scores**: Predicted foreground scores for \ anchors. Its shape is :math:`(N, H W A, 2)`. * **rois**: A bounding box array containing coordinates of \ proposal boxes. This is a concatenation of bounding box \ arrays from multiple images in the batch. \ Its shape is :math:`(R', 4)`. Given :math:`R_i` predicted \ bounding boxes from the :math:`i` th image, \ :math:`R' = \\sum _{i=1} ^ N R_i`. * **roi_indices**: An array containing indices of images to \ which RoIs correspond to. Its shape is :math:`(R',)`. * **anchor**: Coordinates of enumerated shifted anchors. \ Its shape is :math:`(H W A, 4)`. """ n, _, hh, ww = x.shape anchor = _enumerate_shifted_anchor(self.anchor_base, self.feat_stride, hh, ww) anchor = jt.array(anchor) n_anchor = anchor.shape[0] // (hh * ww) h = nn.relu(self.conv1(x)) rpn_locs = self.loc(h) rpn_locs = rpn_locs.permute(0, 2, 3, 1).view(n, -1, 4) rpn_scores = self.score(h) rpn_scores = rpn_scores.permute(0, 2, 3, 1) rpn_softmax_scores = nn.softmax(rpn_scores.view( n, hh, ww, n_anchor, 2), dim=4) rpn_fg_scores = rpn_softmax_scores[:, :, :, :, 1] rpn_fg_scores = rpn_fg_scores.view(n, -1) rpn_scores = rpn_scores.view(n, -1, 2) rois = [] roi_indices = [] for i in range(n): roi = self.proposal_layer(rpn_locs[i], rpn_fg_scores[i], anchor, img_size, scale) batch_index = i * jt.ones((len(roi), ), dtype='int32') rois.append(roi) roi_indices.append(batch_index) rois = jt.contrib.concat(rois, dim=0) roi_indices = jt.contrib.concat(roi_indices, dim=0) return rpn_locs, rpn_scores, rois, roi_indices, anchor
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 __getitem__(self, index): AB_path = self.files[index % len(self.files)] img = Image.open(AB_path) w, h = img.size img_A = img.crop((0, 0, w / 2, h)) img_B = img.crop((w / 2, 0, w, h)) flip = random.random() > 0.5 params = {'load_h': self.load_h, 'load_w': self.load_w, 'flip': flip} transform_A = get_transform(params) transform_B = get_transform(params, gray=True) transform_mask = get_transform(params, gray=True, mask=True) item_A = transform_A(img_A) item_A = jt.array(item_A) item_B = transform_B(img_B) item_B = jt.array(item_B) item_A_l = {} item_B_l = {} regions = ['eyel','eyer','nose','mouth'] basen = os.path.basename(AB_path)[:-4] lm_path = os.path.join(self.lmdir, basen+'.txt') feats = getfeats(lm_path) if flip: for i in range(5): feats[i,0] = self.load_w - feats[i,0] - 1 tmp = [feats[0,0],feats[0,1]] feats[0,:] = [feats[1,0],feats[1,1]] feats[1,:] = tmp mouth_x = int((feats[3,0]+feats[4,0])/2.0) mouth_y = int((feats[3,1]+feats[4,1])/2.0) ratio = self.load_h // 256 rhs = np.array([EYE_H,EYE_H,NOSE_H,MOUTH_H]) * ratio rws = np.array([EYE_W,EYE_W,NOSE_W,MOUTH_W]) * ratio center = np.array([[feats[0,0],feats[0,1]-4*ratio],[feats[1,0],feats[1,1]-4*ratio],[feats[2,0],feats[2,1]-rhs[2]//2+16*ratio],[mouth_x,mouth_y]]) soft_border_mask4 = [] for i in range(4): xb = [np.zeros(rhs[i]),np.ones(rhs[i])*(rws[i]-1)] yb = [np.zeros(rws[i]),np.ones(rws[i])*(rhs[i]-1)] soft_border_mask = getSoft([rhs[i],rws[i]],xb,yb) soft_border_mask = jt.array(soft_border_mask).unsqueeze(0).float() soft_border_mask4.append(soft_border_mask) for i in range(4): item_A_l[regions[i]+'_A'] = item_A[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] * soft_border_mask4[i].repeat(3,1,1) item_B_l[regions[i]+'_B'] = item_B[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] * soft_border_mask4[i] cmasks = [] for i in range(4): if not flip or i not in [2,3]: cmaskpath = os.path.join(self.cmaskdir.format(regions[i]),basen+'.png') else: cmaskpath = os.path.join(self.cmaskdir.format(regions[1-i]),basen+'.png') im_cmask = Image.open(cmaskpath) cmask0 = transform_mask(im_cmask) cmask0 = jt.array(cmask0) cmask0 = (cmask0 >= 0.5).float() cmask = cmask0.clone() cmask = cmask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] cmasks.append(cmask) mask = jt.ones([1,item_A.shape[1],item_A.shape[2]]) # mask out eyes, nose, mouth for i in range(4): mask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] = 0 imgsize = self.load_h maskn = mask[0].numpy() masks = [np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize])] masks[0][1:] = maskn[:-1] masks[1][:-1] = maskn[1:] masks[2][:,1:] = maskn[:,:-1] masks[3][:,:-1] = maskn[:,1:] masks2 = [maskn-e for e in masks] bound = np.minimum.reduce(masks2) bound = -bound xb = [] yb = [] for i in range(4): xbi = [center[i,0]-rws[i]/2, center[i,0]+rws[i]/2-1] ybi = [center[i,1]-rhs[i]/2, center[i,1]+rhs[i]/2-1] for j in range(2): maskx = bound[:,int(xbi[j])] masky = bound[int(ybi[j]),:] tmp_a = maskx * xbi[j] tmp_b = 1-maskx xb += [tmp_b*10000 + tmp_a] tmp_a = masky * ybi[j] tmp_b = 1-masky yb += [tmp_b*10000 + tmp_a] soft = 1-getSoft([imgsize,imgsize],xb,yb) soft = jt.array(soft).unsqueeze(0).float() mask = (jt.ones(mask.shape)-mask)*soft + mask bgpath = os.path.join(self.maskbgdir, basen+'.png') im_bg = Image.open(bgpath) mask2 = transform_mask(im_bg) # mask out background mask2 = jt.array(mask2) mask2 = (mask2 >= 0.5).float() # foreground: 1, background: 0 item_A_l['hair_A'] = (item_A/2+0.5) * mask.repeat(3,1,1) * mask2.repeat(3,1,1) * 2 - 1 item_A_l['bg_A'] = (item_A/2+0.5) * (jt.ones(mask2.shape)-mask2).repeat(3,1,1) * 2 - 1 item_B_l['hair_B'] = (item_B/2+0.5) * mask * mask2 * 2 - 1 item_B_l['bg_B'] = (item_B/2+0.5) * (jt.ones(mask2.shape)-mask2) * 2 - 1 facepath = os.path.join(self.maskfacedir, basen+'.png') im_face = Image.open(facepath) maskface = transform_mask(im_face) maskface = jt.array(maskface) img = tocv2(item_B) dt1, dt2 = dt(img) dt1 = jt.array(dt1) dt2 = jt.array(dt2) dt1 = dt1.unsqueeze(0) dt2 = dt2.unsqueeze(0) return item_A, item_A_l['eyel_A'], item_A_l['eyer_A'], item_A_l['nose_A'], item_A_l['mouth_A'], item_A_l['hair_A'], item_A_l['bg_A'], item_B, item_B_l['eyel_B'], item_B_l['eyer_B'], item_B_l['nose_B'], item_B_l['mouth_B'], item_B_l['hair_B'], item_B_l['bg_B'], mask, mask2, center, dt1, dt2, cmasks[0], cmasks[1], cmasks[2], cmasks[3], maskface
def test_getitem(self): # test for different slice type arr0 = jt.random((4, 3)) arr0_res = arr0[2, :] arr0_res.data[1] = 1 assert arr0[2, 1] == 1 arr1 = jt.array([1, 2, 3, 4]) arr1_res = arr1[None] arr1_res.data[0, 2] = -1 assert arr1[2] == -1 arr2 = jt.array([1, 2, 3, 4]) arr2_res = arr2[...] arr2_res.data[2] = -1 assert arr2[2] == -1 arr3 = jt.array([1, 2, 3, 4]) arr3_res = arr3[3] arr3_res.data[0] = -1 assert arr3[3] == -1 arr4 = jt.random((4, 2, 3, 3)) arr4_res = arr4[..., :, :] arr4_res.data[0, 0, 1, 1] = 1 assert arr4[0, 0, 1, 1] == 1 arr4 = jt.random((4, 2, 3, 3)) arr4_res = arr4[..., :, :2] arr4_res.data[0, 0, 1, 1] = 1 assert arr4[0, 0, 1, 1] != 1 arr4 = jt.random((3, 3)) arr4_res = arr4[..., :, :2] arr4_res.data[1, 1] = 1 assert arr4[1, 1] != 1 arr5 = jt.random((4, 2, 3, 3)) arr5_res = arr5[1:3, :, :, :] arr5_res.data[1, 0, 1, 1] = 1 assert arr5[2, 0, 1, 1] == 1 arr6 = jt.random((4, 2, 3, 3)) arr6_res = arr6[1] arr6_res.data[0, 1, 1] = 1 assert arr6[1, 0, 1, 1] == 1 # test for different data type (float32/float64/bool/int8/int32) arr_float32 = jt.random((4, 2, 3)) arr_float32_res = arr_float32[1:3, :, :] arr_float32_res.data[0, 0, 0] = 1 assert arr_float32[1, 0, 0] == 1 arr_float32_res.data[1, 1, 2] = 1 assert arr_float32[2, 1, 2] == 1 arr_float32[1, 0, 0] = 0 # getitem and setitem do not conflict assert arr_float32_res[0, 0, 0] == 1 arr_bool = jt.bool(np.ones((4, 2, 3))) arr_bool_res = arr_bool[1:3, :, :] arr_bool_res.data[0, 0, 0] = False assert arr_bool[1, 0, 0] == False arr_bool_res.data[0, 0, 1] = False assert arr_bool[1, 0, 1] == False arr_float64 = jt.random((4, 2, 3), dtype='float64') arr_float64_res = arr_float64[1:3, :, :] arr_float64_res.data[0, 0, 0] = 1 assert arr_float64[1, 0, 0] == 1 arr_float64_res.data[1, 1, 2] = 1 assert arr_float64[2, 1, 2] == 1 arr_int32 = jt.ones((4, 2, 3), dtype='int32') arr_int32_res = arr_int32[1:3, :, :] arr_int32_res.data[0, 0, 0] = 0 assert arr_int32[1, 0, 0] == 0 arr_int32_res.data[1, 1, 2] = 0 assert arr_int32[2, 1, 2] == 0
def add_with_mask(A, B, mask): return ((A/2+0.5)*mask + (B/2+0.5)*(jt.ones(mask.shape)-mask))*2-1
def __init__(self, vertices, faces, textures=None, texture_res=1, texture_type='surface', dr_type='softras', metallic_textures=None, roughness_textures=None): ''' vertices, faces and textures (if not None) are expected to be Tensor objects ''' self._vertices = vertices self._faces = faces if isinstance(self._vertices, np.ndarray): self._vertices = jt.array(self._vertices).float() if isinstance(self._faces, np.ndarray): self._faces = jt.array(self._faces).int() if len(self._vertices.shape) == 2: self._vertices = self._vertices.unsqueeze(0) if len(self._faces.shape) == 2: self._faces = self._faces.unsqueeze(0) self.texture_type = texture_type self.batch_size = self._vertices.shape[0] self.num_vertices = self._vertices.shape[1] self.num_faces = self._faces.shape[1] self._face_vertices = None self._face_vertices_update = True self._surface_normals = None self._surface_normals_update = True self._vertex_normals = None self._vertex_normals_update = True self._with_specular = True self._fill_back = False self.dr_type = dr_type if texture_type == 'surface': if self.dr_type == 'softras': self._metallic_textures = jt.zeros( (self.batch_size, self.num_faces, texture_res**2, 1)) self._roughness_textures = jt.ones( (self.batch_size, self.num_faces, texture_res**2, 1)) elif self.dr_type == 'n3mr': self._metallic_textures = jt.zeros( (self.batch_size, self.num_faces, texture_res, texture_res, texture_res, 1)) self._roughness_textures = jt.ones( (self.batch_size, self.num_faces, texture_res, texture_res, texture_res, 1)) elif texture_type == 'vertex': self._metallic_textures = jt.zeros( (self.batch_size, self.num_vertices, 1)) self._roughness_textures = jt.ones( (self.batch_size, self.num_vertices, 1)) if metallic_textures is not None: self._metallic_textures = metallic_textures if roughness_textures is not None: self._roughness_textures = roughness_textures # create textures if textures is None: if texture_type == 'surface': if self.dr_type == 'softras': self._textures = jt.ones( (self.batch_size, self.num_faces, texture_res**2, 3)) elif self.dr_type == 'n3mr': self._textures = jt.ones( (self.batch_size, self.num_faces, texture_res, texture_res, texture_res, 3)) self.texture_res = texture_res elif texture_type == 'vertex': self._textures = jt.ones( (self.batch_size, self.num_vertices, 3)) self.texture_res = 1 else: if isinstance(textures, np.ndarray): textures = jt.array(textures).float() if len(textures.shape) == 3 and texture_type == 'surface': textures = textures.unsqueeze(0) if len(textures.shape) == 2 and texture_type == 'vertex': textures = textures.unsqueeze(0) if len(textures.shape) == 5: textures = textures.unsqueeze(0) self._textures = textures self.texture_res = int(np.sqrt(self._textures.shape[2])) self._origin_vertices = self._vertices self._origin_faces = self._faces self._origin_textures = self._textures
def __init__(self, n): super(FrozenBatchNorm2d, self).__init__() self.weight = jt.ones(n).stop_grad() self.bias = jt.zeros(n).stop_grad() self.running_mean = jt.zeros(n).stop_grad() self.running_var = jt.ones(n).stop_grad()
def test_flatten(self): a = jt.ones([2,3,4]) assert a.flatten().shape == [24] assert a.flatten(1).shape == [2,12] assert a.flatten(0,-2).shape == [6,4]
def test_tensor_bad_types_to_pil_image(self): with self.assertRaises(ValueError): transform.ToPILImage()(jt.ones((1, 3, 4, 4)))
# Sample noise z = jt.array(np.random.normal( 0, 1, (n_row**2, opt.latent_dim))).float32().stop_grad() gen_imgs = decoder(z) save_image(gen_imgs.numpy(), "images/%d.png" % batches_done, nrow=n_row) # ---------- # Training # ---------- for epoch in range(opt.n_epochs): for i, (imgs, _) in enumerate(train_loader): sta = time.time() # Adversarial ground truths valid = jt.ones([imgs.shape[0], 1]).stop_grad() fake = jt.zeros([imgs.shape[0], 1]).stop_grad() # Configure input real_imgs = jt.array(imgs).stop_grad() # ----------------- # Train Generator # ----------------- encoded_imgs = encoder(real_imgs) decoded_imgs = decoder(encoded_imgs) # Loss measures generator's ability to fool the discriminator g_loss = ( 0.001 * adversarial_loss(discriminator(encoded_imgs), valid) + 0.999 * pixelwise_loss(decoded_imgs, real_imgs))
nrow=n_row) save_image(sample2.numpy(), "images/varying_c2/%d.png" % batches_done, nrow=n_row) # ---------- # Training # ---------- for epoch in range(opt.n_epochs): for i, (imgs, labels) in enumerate(dataloader): batch_size = imgs.shape[0] # Adversarial ground truths valid = jt.ones((batch_size, 1)).float32().stop_grad() fake = jt.zeros((batch_size, 1)).float32().stop_grad() # Configure input real_imgs = jt.array(imgs).float32() labels = to_categorical(labels.numpy(), num_columns=opt.n_classes) # ----------------- # Train Generator # ----------------- # Sample noise and labels as generator input z = jt.array(np.random.normal(0, 1, (batch_size, opt.latent_dim))).float32() label_input = to_categorical(np.random.randint(0, opt.n_classes, batch_size),
loss = jt.mean(mu_2) return loss # ---------- # Training # ---------- prev_time = time.time() for epoch in range(opt.epoch, opt.n_epochs): for i, batch in enumerate(dataloader): # Set model input X1 = batch[0].stop_grad() X2 = batch[1].stop_grad() # Adversarial ground truths valid = jt.ones((X1.size(0), *D1.output_shape)).stop_grad() fake = jt.zeros((X1.size(0), *D1.output_shape)).stop_grad() # ------------------------------- # Train Encoders and Generators # ------------------------------- E1.train() E2.train() G1.train() G2.train() # Get shared latent representation mu1, Z1 = E1(X1) mu2, Z2 = E2(X2) # Reconstruct images
def test_state_dict(self): a = jt.ones(2) opt = jt.optim.SGD([a], 0.1) s = opt.state_dict() # print(s) opt.load_state_dict(s)
def execute(self, x): """ The input should be of size [batch_size, 3, img_h, img_w] """ _, _, img_h, img_w = x.shape cfg._tmp_img_h = img_h cfg._tmp_img_w = img_w with timer.env('backbone'): outs = self.backbone(x) if cfg.fpn is not None: with timer.env('fpn'): # Use backbone.selected_layers because we overwrote self.selected_layers outs = [outs[i] for i in cfg.backbone.selected_layers] outs = self.fpn(outs) proto_out = None if cfg.mask_type == mask_type.lincomb and cfg.eval_mask_branch: with timer.env('proto'): proto_x = x if self.proto_src is None else outs[self.proto_src] if self.num_grids > 0: grids = self.grid.repeat(proto_x.shape[0], 1, 1, 1) proto_x = jt.contrib.concat([proto_x, grids], dim=1) proto_out = self.proto_net(proto_x) proto_out = cfg.mask_proto_prototype_activation(proto_out) if cfg.mask_proto_prototypes_as_features: # Clone here because we don't want to permute this, though idk if contiguous makes this unnecessary proto_downsampled = proto_out.clone() if cfg.mask_proto_prototypes_as_features_no_grad: proto_downsampled = proto_out.detach() # Move the features last so the multiplication is easy proto_out = proto_out.permute(0, 2, 3, 1) if cfg.mask_proto_bias: bias_shape = [x for x in proto_out.shape] bias_shape[-1] = 1 proto_out = jt.contrib.concat( [proto_out, jt.ones(bias_shape)], -1) with timer.env('pred_heads'): pred_outs = {'loc': [], 'conf': [], 'mask': [], 'priors': []} if cfg.use_mask_scoring: pred_outs['score'] = [] if cfg.use_instance_coeff: pred_outs['inst'] = [] for idx, pred_layer in zip(self.selected_layers, self.prediction_layers): pred_x = outs[idx] if cfg.mask_type == mask_type.lincomb and cfg.mask_proto_prototypes_as_features: # Scale the prototypes down to the current prediction layer's size and add it as inputs proto_downsampled = nn.interpolate( proto_downsampled, size=outs[idx].shape[2:], mode='bilinear', align_corners=False) # proto_downsampled = interpolate(proto_downsampled, size=outs[idx].shape[2:], mode='bilinear', align_corners=False) pred_x = jt.contrib.concat([pred_x, proto_downsampled], dim=1) # A hack for the way dataparallel works if cfg.share_prediction_module and pred_layer is not self.prediction_layers[ 0]: pred_layer.parent = [self.prediction_layers[0]] p = pred_layer(pred_x) for k, v in p.items(): pred_outs[k].append(v) for k, v in pred_outs.items(): pred_outs[k] = jt.contrib.concat(v, -2) if proto_out is not None: pred_outs['proto'] = proto_out #print('hh',pred_outs) #print() if self.is_training(): # For the extra loss functions if cfg.use_class_existence_loss: pred_outs['classes'] = self.class_existence_fc( outs[-1].mean(dim=(2, 3))) if cfg.use_semantic_segmentation_loss: pred_outs['segm'] = self.semantic_seg_conv(outs[0]) return pred_outs else: if cfg.use_mask_scoring: pred_outs['score'] = jt.sigmoid(pred_outs['score']) if cfg.use_focal_loss: if cfg.use_sigmoid_focal_loss: # Note: even though conf[0] exists, this mode doesn't train it so don't use it pred_outs['conf'] = jt.sigmoid(pred_outs['conf']) if cfg.use_mask_scoring: pred_outs['conf'] *= pred_outs['score'] elif cfg.use_objectness_score: # See focal_loss_sigmoid in multibox_loss.py for details objectness = jt.sigmoid(pred_outs['conf'][:, :, 0]) pred_outs['conf'][:, :, 1:] = objectness.unsqueeze( 2) * nn.softmax(pred_outs['conf'][:, :, 1:], -1) pred_outs['conf'][:, :, 0] = 1 - objectness else: pred_outs['conf'] = nn.softmax(pred_outs['conf'], -1) else: if cfg.use_objectness_score: objectness = jt.sigmoid(pred_outs['conf'][:, :, 0]) pred_outs['conf'][:, :, 1:] = (objectness > 0.10).unsqueeze(-1) \ * nn.softmax(pred_outs['conf'][:, :, 1:], dim=-1) else: pred_outs['conf'] = nn.softmax(pred_outs['conf'], -1) return self.detect(pred_outs, self)
def test_opt_grad(self): a = jt.ones(2) opt = jt.optim.SGD([a], 0.1) opt.backward(a**2) g = a.opt_grad(opt) np.testing.assert_allclose(g.data, 2)
warmup_times = -1 run_times = 3000 total_time = 0. cnt = 0 # ---------- # Training # ---------- prev_time = time.time() for epoch in range(opt.epoch, opt.n_epochs): for i, (real_B, real_A) in enumerate(dataloader): # Adversarial ground truths valid = jt.ones([real_A.shape[0], 1]).stop_grad() fake = jt.zeros([real_A.shape[0], 1]).stop_grad() # ------------------ # Train Generators # ------------------ # GAN loss fake_B = generator(real_A) pred_fake = discriminator(fake_B, real_A) loss_GAN = criterion_GAN(pred_fake, valid) # Pixel-wise loss loss_pixel = criterion_pixelwise(fake_B, real_B) # Total loss loss_G = loss_GAN + lambda_pixel * loss_pixel optimizer_G.step(loss_G)
def __getitem__(self, index): AB_path = self.files[index % len(self.files)] img = Image.open(AB_path) w, h = img.size img_A = img.crop((0, 0, w / 2, h)) img_B = img.crop((w / 2, 0, w, h)) flip = random.random() > 0.5 params = {'load_h': self.load_h, 'load_w': self.load_w, 'flip': flip} transform_A = get_transform(params) transform_B = get_transform(params, gray=True) transform_mask = get_transform(params, gray=True, mask=True) item_A = transform_A(img_A) item_A = jt.array(item_A) item_B = transform_B(img_B) item_B = jt.array(item_B) item_A_l = {} regions = ['eyel', 'eyer', 'nose', 'mouth'] basen = os.path.basename(AB_path)[:-4] lm_path = os.path.join(self.lmdir, basen + '.txt') feats = getfeats(lm_path) if flip: for i in range(5): feats[i, 0] = self.load_w - feats[i, 0] - 1 tmp = [feats[0, 0], feats[0, 1]] feats[0, :] = [feats[1, 0], feats[1, 1]] feats[1, :] = tmp mouth_x = int((feats[3, 0] + feats[4, 0]) / 2.0) mouth_y = int((feats[3, 1] + feats[4, 1]) / 2.0) ratio = self.load_h // 256 rhs = np.array([EYE_H, EYE_H, NOSE_H, MOUTH_H]) * ratio rws = np.array([EYE_W, EYE_W, NOSE_W, MOUTH_W]) * ratio center = np.array( [[feats[0, 0], feats[0, 1] - 4 * ratio], [feats[1, 0], feats[1, 1] - 4 * ratio], [feats[2, 0], feats[2, 1] - rhs[2] // 2 + 16 * ratio], [mouth_x, mouth_y]]) for i in range(4): item_A_l[regions[i] + '_A'] = item_A[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2), int(center[i, 0] - rws[i] / 2):int(center[i, 0] + rws[i] / 2)] mask = jt.ones([1, item_A.shape[1], item_A.shape[2]]) # mask out eyes, nose, mouth for i in range(4): mask[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2), int(center[i, 0] - rws[i] / 2):int(center[i, 0] + rws[i] / 2)] = 0 bgpath = os.path.join(self.maskdir, basen + '.png') im_bg = Image.open(bgpath) mask2 = transform_mask(im_bg) # mask out background mask2 = jt.array(mask2) mask2 = (mask2 >= 0.5).float() # foreground: 1, background: 0 item_A_l['hair_A'] = (item_A / 2 + 0.5) * mask.repeat( 3, 1, 1) * mask2.repeat(3, 1, 1) * 2 - 1 item_A_l['bg_A'] = (item_A / 2 + 0.5) * (jt.ones(mask2.shape) - mask2).repeat(3, 1, 1) * 2 - 1 img = tocv2(item_B) dt1, dt2 = dt(img) dt1 = jt.array(dt1) dt2 = jt.array(dt2) dt1 = dt1.unsqueeze(0) dt2 = dt2.unsqueeze(0) return item_A, item_A_l['eyel_A'], item_A_l['eyer_A'], item_A_l[ 'nose_A'], item_A_l['mouth_A'], item_A_l['hair_A'], item_A_l[ 'bg_A'], mask, mask2, center, item_B, dt1, dt2