jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_bboxes_logi(self,cls_scores, mask_preds, img_metas, cfg, rescale=None, score_thr=True):
		_, _, b_h, b_w = mask_preds[0].shape
		crop_h, crop_w, _ = img_metas[0]['img_shape']
		ori_h, ori_w, _ = img_metas[0]['ori_shape']
		cls_scores_ori = []
		for i in range(5):
			cls_sc = cls_scores[i].clone().detach()
			cls_sc = F.interpolate(cls_sc, (ori_h, ori_w), mode='bilinear')
			cls_scores_ori.append(cls_sc)
		# cat dim: default 0
		if self.use_sigmoid:
			flatten_cls_scores = torch.cat([cls_score.permute(0, 2, 3, 1).reshape(-1, self.cls_out_channels) for cls_score in cls_scores]).sigmoid()
		else:
			flatten_cls_scores = torch.cat([cls_score.permute(0, 2, 3, 1).reshape(-1, self.cls_out_channels) for cls_score in cls_scores]).softmax(dim=1)
			flatten_cls_scores = flatten_cls_scores[:,1:]  
		for i in range(5):
			mask_preds[i] = F.interpolate(mask_preds[i],(b_h,b_w), mode='bilinear')
		mask_preds = torch.cat(mask_preds,dim=1)[0]
		# logit and class (per pixel)
		scores, labels = torch.max(flatten_cls_scores,dim=-1)
		nms_pre = cfg.get('nms_pre', -1)
		if score_thr:
			score_thr = cfg.get('score_thr', -1)
		else:
			score_thr = 0
		mask_thr = cfg.get('mask_thr_binary',-1)
		det_masks = np.array([])
		det_bboxes = np.zeros((0,5))
		det_labels = np.zeros(0).astype(int)
		# store feature scale
		sc_ind = []
		for i in range(5):
			sc_ind.append(torch.zeros((self.grid_num[i] ** 2, 1)) + i)
		sc_ind = torch.cat(sc_ind, axis=0)
		if nms_pre > 0 and len(scores) >=1 and scores.max() > score_thr:
			# HERE: need to change scores for softmax
			valid_inds = torch.nonzero(scores >= score_thr).squeeze()
			if valid_inds.dim()!=0:
				scores = scores[valid_inds]
				labels = labels[valid_inds]
				mask_preds = mask_preds[valid_inds]
				sc_ind = sc_ind[valid_inds]

				if scores.shape[0] < nms_pre:
					nms_pre = scores.shape[0]

				_, topk_inds = scores.topk(nms_pre)
				scores = scores[topk_inds]
				labels = labels[topk_inds]
				mask_preds = mask_preds[topk_inds]
				sc_ind = sc_ind[topk_inds]
				mask_preds = F.interpolate(mask_preds.unsqueeze(0), (b_h*self.strides[0], b_w*self.strides[0]), mode='bilinear')
				mask_preds = mask_preds[:, :, :crop_h, :crop_w]
				mask_preds = F.sigmoid(F.interpolate(mask_preds, (ori_h, ori_w), mode='bilinear'))[0]
				mask_preds = mask_preds > mask_thr

				masks = self.nms(scores, labels, mask_preds, sc_ind, cfg.nms.iou_thr)
				n = len(masks)
				det_masks = []
				det_bboxes = np.zeros((n, 5))
				det_labels = np.zeros(n).astype(int)
				det_scs = np.zeros(n).astype(int)

				for i in range(n):
					det_bboxes[i, -1] = masks[i][-2]
					det_labels[i] = masks[i][-3]
					det_masks.append(masks[i][0])
					det_scs[i] = masks[i][-1]
				det_masks = np.array(det_masks)

				if self.out_path is not None:
					self.infer_vis(img_metas, cls_scores, det_masks, det_bboxes, det_labels, det_scs, out_path=self.out_path)

		return det_bboxes, det_labels, det_masks
//  [python]
def upsample_bilinear2d_python(node, speedup):
    class UpsampleModule(torch.nn.Module):
        def __init__(self, size_list, scale_list):
            super(UpsampleModule, self).__init__()
            self.size_list = size_list
            self.scale_list = scale_list

        def forward(self, *args):
            """
            The first input of args is the target tensor to upsample
            , the following parameters is useless, because we already
            get the size_list and the scale_list by parsing the cpp_nodes.
            """
            return torch.nn.functional.upsample_bilinear(args[0],
                                                         size=self.size_list, scale_factor=self.scale_list)
    c_node = node.key_node
    inputs = list(c_node.inputs())
    size_list_node = inputs[1].node()
    scale_list_node = inputs[3].node()
    size_list = None
    scale_list = None
    
    if size_list_node.kind() == 'prim::ListConstruct':
        size_list = translate_list(inputs[1], speedup)
    if scale_list_node.kind() == 'prim::ListConstruct':
        scale_list = translate_list(inputs[3], speedup)
    return UpsampleModule(size_list, scale_list)
//  [python]
def upsample_bilinear2d_python(node, speedup):
    class UpsampleModule(torch.nn.Module):
        def __init__(self, size_list, scale_list):
            super(UpsampleModule, self).__init__()
            self.size_list = size_list
            self.scale_list = scale_list

        def forward(self, *args):
            """
            The first input of args is the target tensor to upsample
            , the following parameters is useless, because we already
            get the size_list and the scale_list by parsing the cpp_nodes.
            """
            return torch.nn.functional.interpolate(args[0],
                                                         size=self.size_list, scale_factor=self.scale_list, mode='bilinear')
    c_node = node.key_node
    inputs = list(c_node.inputs())
    size_list_node = inputs[1].node()
    scale_list_node = inputs[3].node()
    size_list = None
    scale_list = None
    
    if size_list_node.kind() == 'prim::ListConstruct':
        size_list = translate_list(inputs[1], speedup)
    if scale_list_node.kind() == 'prim::ListConstruct':
        scale_list = translate_list(inputs[3], speedup)
    return UpsampleModule(size_list, scale_list)
//  [python]
def upsample_nearest2d_python(node, speedup):
    class UpsampleModule(torch.nn.Module):
        def __init__(self, size_list, scale_list):
            super(UpsampleModule, self).__init__()
            self.size_list = size_list
            self.scale_list = scale_list

        def forward(self, *args):
            """
            The first input of args is the target tensor to upsample
            , the following parameters is useless, because we already
            get the size_list and the scale_list by parsing the cpp_nodes.
            """
            return torch.nn.functional.upsample_nearest(args[0],
                                                         size=self.size_list, scale_factor=self.scale_list)
    c_node = node.key_node
    inputs = list(c_node.inputs())
    # assert len(inputs)==3
    size_list_node = inputs[1].node()
    scale_list_node = inputs[2].node()
    
    size_list = None
    scale_list = None
    
    if size_list_node.kind() == 'prim::ListConstruct':
        size_list = translate_list(inputs[1], speedup)
        assert isinstance(size_list[0], int)
    if scale_list_node.kind() == 'prim::ListConstruct':
        scale_list = translate_list(inputs[2], speedup)
        assert isinstance(scale_list[0], float)
    return UpsampleModule(size_list, scale_list)
//  [python]
def upsample_nearest2d_python(node, speedup):
    class UpsampleModule(torch.nn.Module):
        def __init__(self, size_list, scale_list):
            super(UpsampleModule, self).__init__()
            self.size_list = size_list
            self.scale_list = scale_list

        def forward(self, *args):
            """
            The first input of args is the target tensor to upsample
            , the following parameters is useless, because we already
            get the size_list and the scale_list by parsing the cpp_nodes.
            """
            return torch.nn.functional.interpolate(args[0],
                                                         size=self.size_list, scale_factor=self.scale_list, mode='nearest')
    c_node = node.key_node
    inputs = list(c_node.inputs())
    # assert len(inputs)==3
    size_list_node = inputs[1].node()
    scale_list_node = inputs[2].node()
    
    size_list = None
    scale_list = None
    
    if size_list_node.kind() == 'prim::ListConstruct':
        size_list = translate_list(inputs[1], speedup)
        assert isinstance(size_list[0], int)
    if scale_list_node.kind() == 'prim::ListConstruct':
        scale_list = translate_list(inputs[2], speedup)
        assert isinstance(scale_list[0], float)
    return UpsampleModule(size_list, scale_list)
//  [python]
def attribution_random_patch_augmentation(
        x,
        cam,
        image_size=84,
        output_size=4,
        quantile=0.90,
        patch_proba=0.7,
        dataset="places365_standard",
):
    """from SGQN"""
    if dataset == "places365_standard":
        if places_dataloader is None:
            _load_places(batch_size=x.size(0), image_size=x.size(-1))
        negative = _get_places_batch(batch_size=x.size(0)).repeat(
            1, x.size(1) // 3, 1, 1
        )
    else:
        raise NotImplementedError(
            f'overlay has not been implemented for dataset "{dataset}"'
        )
    cam = cam.to(x.device)
    cam = F.adaptive_avg_pool2d(cam, output_size=output_size)
    q = torch.quantile(cam.flatten(1), quantile, 1)
    mask = (cam >= q[:, None, None]).long()
    exploration_mask = torch.rand(*mask.shape).to(x.device)
    exploration_mask[~mask] = 0
    expl_max = torch.amax(exploration_mask.view(mask.size(0), -1), dim=1)
    exploration_mask = (
            exploration_mask.view(-1, mask.size(1), mask.size(2)) == expl_max[:, None, None]
    ).long()
    bern = torch.bernoulli(torch.ones_like(mask) * patch_proba).long().to(x.device)
    selected_patch = (mask * bern) + exploration_mask
    selected_patch[selected_patch > 1] = 1
    selected_patch = F.upsample_nearest(selected_patch.float().unsqueeze(1), image_size)
    # augmented_x = (((0.5) * (x / 255.0) + (0.5) * negative) * 255.0) * selected_patch
    augmented_x = x * selected_patch
    complementary_mask = ~(selected_patch.bool())
    negative = negative * (complementary_mask.float())
    return augmented_x + (negative * 255)
//  [python]
def attribution_random_patch_augmentation(
        x,
        cam,
        image_size=84,
        output_size=4,
        quantile=0.90,
        patch_proba=0.7,
        dataset="places365_standard",
):
    """from SGQN"""
    if dataset == "places365_standard":
        if places_dataloader is None:
            _load_places(batch_size=x.size(0), image_size=x.size(-1))
        negative = _get_places_batch(batch_size=x.size(0)).repeat(
            1, x.size(1) // 3, 1, 1
        )
    else:
        raise NotImplementedError(
            f'overlay has not been implemented for dataset "{dataset}"'
        )
    cam = cam.to(x.device)
    cam = F.adaptive_avg_pool2d(cam, output_size=output_size)
    q = torch.quantile(cam.flatten(1), quantile, 1)
    mask = (cam >= q[:, None, None]).long()
    exploration_mask = torch.rand(*mask.shape).to(x.device)
    exploration_mask[~mask] = 0
    expl_max = torch.amax(exploration_mask.view(mask.size(0), -1), dim=1)
    exploration_mask = (
            exploration_mask.view(-1, mask.size(1), mask.size(2)) == expl_max[:, None, None]
    ).long()
    bern = torch.bernoulli(torch.ones_like(mask) * patch_proba).long().to(x.device)
    selected_patch = (mask * bern) + exploration_mask
    selected_patch[selected_patch > 1] = 1
    selected_patch = F.interpolate(selected_patch.float().unsqueeze(1), size=image_size, mode='nearest')
    # augmented_x = (((0.5) * (x / 255.0) + (0.5) * negative) * 255.0) * selected_patch
    augmented_x = x * selected_patch
    complementary_mask = ~(selected_patch.bool())
    negative = negative * (complementary_mask.float())
    return augmented_x + (negative * 255)
//  [python]
def scale2(x):
    if x.size(2) > 128: # do not need to scale the input
        return x
    x = torch.nn.functional.upsample(x, scale_factor=2, mode='nearest')  #bicubic is not available for the time being.
    return x
//  [python]
def scale2(x):
    if x.size(2) > 128: # do not need to scale the input
        return x
    x = torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')  #bicubic is not available for the time being.
    return x
//  [python]
def OverlayBeliefOnImage(img, beliefs, name="tmp", path="", factor=0.7, grid=3, 
        norm_belief = True, save = False, scale_factor=1):
    """
    take as input 
    img: a tensor image in pytorch normalized at 0.5
            3xwxh
    belief: tensor of the same size as the image to overlay over img 
            nb_beliefxwxh
    name: str to name the image, e.g., output.png
    path: where to save, e.g., /where/to/save/
    factor: float [0,1] how much to keep the original, 1 = fully, 0 black
    grid: how big the grid, e.g., 3 wide. 
    norm_belief: bool to normalize the values [0,1]
    save: write to disk 
    scale_factor: how much to scale the belief map
    returns the image 
    """

    belief_imgs = []
    if not img is None:
        in_img = img
        in_img *= factor     
    tensor = nn.functional.upsample(beliefs.unsqueeze(0),scale_factor=scale_factor ).squeeze()      
    for j in range(tensor.size()[0]):
        belief = tensor[j].clone()
        if norm_belief:
            belief -= float(torch.min(belief).data.cpu().numpy())
            belief /= float(torch.max(belief).data.cpu().numpy())
        belief = torch.clamp(belief,0,1).cpu()
        
        if img is None:
            belief = torch.cat([
                        belief.unsqueeze(0),belief.unsqueeze(0),belief.unsqueeze(0)
                        ]).unsqueeze(0)
        else:
            belief = torch.cat([
                        belief.unsqueeze(0)* in_img[0,:,:] + in_img[0,:,:],
                        belief.unsqueeze(0)* in_img[1,:,:] + in_img[1,:,:],
                        belief.unsqueeze(0)* in_img[2,:,:] + in_img[2,:,:]
                        ]).unsqueeze(0)
            
        # print (in_img[0,:,:].shape)
        # print ((belief.unsqueeze(0) * in_img[1,:,:] * factor + in_img[1,:,:]).shape)
        # print (in_img[2,:,:].shape)
        # belief = torch.cat([
        #             in_img[0,:,:].unsqueeze(0),
        #             belief.unsqueeze(0) * in_img[1,:,:] + in_img[1,:,:],
        #             in_img[2,:,:].unsqueeze(0)
        #             ]).unsqueeze(0)
        belief = torch.clamp(belief,0,1) 

        belief_imgs.append(belief.data.squeeze().numpy())

    # Create the image grid
    belief_imgs = torch.tensor(np.array(belief_imgs))

    img,grid = save_image(belief_imgs, "{}{}".format(path, name), 
        mean=0, std=1, nrow=grid, save=save)

    return img, grid
//  [python]
def OverlayBeliefOnImage(img, beliefs, name="tmp", path="", factor=0.7, grid=3, 
        norm_belief = True, save = False, scale_factor=1):
    """
    take as input 
    img: a tensor image in pytorch normalized at 0.5
            3xwxh
    belief: tensor of the same size as the image to overlay over img 
            nb_beliefxwxh
    name: str to name the image, e.g., output.png
    path: where to save, e.g., /where/to/save/
    factor: float [0,1] how much to keep the original, 1 = fully, 0 black
    grid: how big the grid, e.g., 3 wide. 
    norm_belief: bool to normalize the values [0,1]
    save: write to disk 
    scale_factor: how much to scale the belief map
    returns the image 
    """

    belief_imgs = []
    if not img is None:
        in_img = img
        in_img *= factor     
    tensor = nn.functional.interpolate(beliefs.unsqueeze(0),scale_factor=scale_factor ).squeeze()      
    for j in range(tensor.size()[0]):
        belief = tensor[j].clone()
        if norm_belief:
            belief -= float(torch.min(belief).data.cpu().numpy())
            belief /= float(torch.max(belief).data.cpu().numpy())
        belief = torch.clamp(belief,0,1).cpu()
        
        if img is None:
            belief = torch.cat([
                        belief.unsqueeze(0),belief.unsqueeze(0),belief.unsqueeze(0)
                        ]).unsqueeze(0)
        else:
            belief = torch.cat([
                        belief.unsqueeze(0)* in_img[0,:,:] + in_img[0,:,:],
                        belief.unsqueeze(0)* in_img[1,:,:] + in_img[1,:,:],
                        belief.unsqueeze(0)* in_img[2,:,:] + in_img[2,:,:]
                        ]).unsqueeze(0)
            
        # print (in_img[0,:,:].shape)
        # print ((belief.unsqueeze(0) * in_img[1,:,:] * factor + in_img[1,:,:]).shape)
        # print (in_img[2,:,:].shape)
        # belief = torch.cat([
        #             in_img[0,:,:].unsqueeze(0),
        #             belief.unsqueeze(0) * in_img[1,:,:] + in_img[1,:,:],
        #             in_img[2,:,:].unsqueeze(0)
        #             ]).unsqueeze(0)
        belief = torch.clamp(belief,0,1) 

        belief_imgs.append(belief.data.squeeze().numpy())

    # Create the image grid
    belief_imgs = torch.tensor(np.array(belief_imgs))

    img,grid = save_image(belief_imgs, "{}{}".format(path, name), 
        mean=0, std=1, nrow=grid, save=save)

    return img, grid
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        map0, map1, map2, map3, map4 = self.fpn(x)

        map4 = nn.functional.upsample(
            self.head4(map4), scale_factor=8, mode='nearest')
        map3 = nn.functional.upsample(
            self.head3(map3), scale_factor=4, mode='nearest')
        map2 = nn.functional.upsample(
            self.head2(map2), scale_factor=2, mode='nearest')
        map1 = nn.functional.upsample(
            self.head1(map1), scale_factor=1, mode='nearest')

        smoothed = self.smooth(torch.cat([map4, map3, map2, map1], dim=1))
        smoothed = nn.functional.upsample(
            smoothed, scale_factor=2, mode='nearest')
        smoothed = self.smooth2(smoothed + map0)
        smoothed = nn.functional.upsample(
            smoothed, scale_factor=2, mode='nearest')

        final = self.final(smoothed)
        res = torch.tanh(final) + x

        return torch.clamp(res, min=-1, max=1)
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        map0, map1, map2, map3, map4 = self.fpn(x)

        map4 = nn.functional.interpolate(
            self.head4(map4), scale_factor=8, mode='nearest')
        map3 = nn.functional.interpolate(
            self.head3(map3), scale_factor=4, mode='nearest')
        map2 = nn.functional.interpolate(
            self.head2(map2), scale_factor=2, mode='nearest')
        map1 = nn.functional.interpolate(
            self.head1(map1), scale_factor=1, mode='nearest')

        smoothed = self.smooth(torch.cat([map4, map3, map2, map1], dim=1))
        smoothed = nn.functional.interpolate(
            smoothed, scale_factor=2, mode='nearest')
        smoothed = self.smooth2(smoothed + map0)
        smoothed = nn.functional.interpolate(
            smoothed, scale_factor=2, mode='nearest')

        final = self.final(smoothed)
        res = torch.tanh(final) + x

        return torch.clamp(res, min=-1, max=1)
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        # Bottom-up pathway, from ResNet
        enc0 = self.enc0(x)

        enc1 = self.enc1(enc0)  # 256

        enc2 = self.enc2(enc1)  # 512

        enc3 = self.enc3(enc2)  # 1024

        enc4 = self.enc4(enc3)  # 2048

        # Lateral connections

        lateral4 = self.pad(self.lateral4(enc4))
        lateral3 = self.pad(self.lateral3(enc3))
        lateral2 = self.lateral2(enc2)
        lateral1 = self.pad(self.lateral1(enc1))
        lateral0 = self.lateral0(enc0)

        # Top-down pathway
        pad = (1, 2, 1, 2)  # pad last dim by 1 on each side
        pad1 = (0, 1, 0, 1)
        map4 = lateral4
        map3 = self.td1(
            lateral3 +
            nn.functional.upsample(map4, scale_factor=2, mode='nearest'))
        map2 = self.td2(
            F.pad(lateral2, pad, 'reflect') +
            nn.functional.upsample(map3, scale_factor=2, mode='nearest'))
        map1 = self.td3(
            lateral1 +
            nn.functional.upsample(map2, scale_factor=2, mode='nearest'))
        return F.pad(lateral0, pad1, 'reflect'), map1, map2, map3, map4
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        # Bottom-up pathway, from ResNet
        enc0 = self.enc0(x)

        enc1 = self.enc1(enc0)  # 256

        enc2 = self.enc2(enc1)  # 512

        enc3 = self.enc3(enc2)  # 1024

        enc4 = self.enc4(enc3)  # 2048

        # Lateral connections

        lateral4 = self.pad(self.lateral4(enc4))
        lateral3 = self.pad(self.lateral3(enc3))
        lateral2 = self.lateral2(enc2)
        lateral1 = self.pad(self.lateral1(enc1))
        lateral0 = self.lateral0(enc0)

        # Top-down pathway
        pad = (1, 2, 1, 2)  # pad last dim by 1 on each side
        pad1 = (0, 1, 0, 1)
        map4 = lateral4
        map3 = self.td1(
            lateral3 +
            nn.functional.interpolate(map4, scale_factor=2, mode='nearest'))
        map2 = self.td2(
            F.pad(lateral2, pad, 'reflect') +
            nn.functional.interpolate(map3, scale_factor=2, mode='nearest'))
        map1 = self.td3(
            lateral1 +
            nn.functional.interpolate(map2, scale_factor=2, mode='nearest'))
        return F.pad(lateral0, pad1, 'reflect'), map1, map2, map3, map4
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        # Bottom-up pathway, from ResNet
        enc0 = self.enc0(x)

        enc1 = self.enc1(enc0)  # 256

        enc2 = self.enc2(enc1)  # 512

        enc3 = self.enc3(enc2)  # 1024

        enc4 = self.enc4(enc3)  # 2048

        # Lateral connections

        lateral4 = self.pad(self.lateral4(enc4))
        lateral3 = self.pad(self.lateral3(enc3))
        lateral2 = self.lateral2(enc2)
        lateral1 = self.pad(self.lateral1(enc1))
        lateral0 = self.lateral0(enc0)

        # Top-down pathway
        pad = (1, 2, 1, 2)  # pad last dim by 1 on each side
        pad1 = (0, 1, 0, 1)
        map4 = lateral4
        map3 = lateral3 + nn.functional.upsample(
            map4, scale_factor=2, mode='nearest')
        map2 = F.pad(lateral2, pad, 'reflect') + nn.functional.upsample(
            map3, scale_factor=2, mode='nearest')
        map1 = lateral1 + nn.functional.upsample(
            map2, scale_factor=2, mode='nearest')
        return F.pad(lateral0, pad1, 'reflect'), map1, map2, map3, map4
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        # Bottom-up pathway, from ResNet
        enc0 = self.enc0(x)

        enc1 = self.enc1(enc0)  # 256

        enc2 = self.enc2(enc1)  # 512

        enc3 = self.enc3(enc2)  # 1024

        enc4 = self.enc4(enc3)  # 2048

        # Lateral connections

        lateral4 = self.pad(self.lateral4(enc4))
        lateral3 = self.pad(self.lateral3(enc3))
        lateral2 = self.lateral2(enc2)
        lateral1 = self.pad(self.lateral1(enc1))
        lateral0 = self.lateral0(enc0)

        # Top-down pathway
        pad = (1, 2, 1, 2)  # pad last dim by 1 on each side
        pad1 = (0, 1, 0, 1)
        map4 = lateral4
        map3 = lateral3 + nn.functional.interpolate(
            map4, scale_factor=2, mode='nearest')
        map2 = F.pad(lateral2, pad, 'reflect') + nn.functional.interpolate(
            map3, scale_factor=2, mode='nearest')
        map1 = lateral1 + nn.functional.interpolate(
            map2, scale_factor=2, mode='nearest')
        return F.pad(lateral0, pad1, 'reflect'), map1, map2, map3, map4
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        # Bottom-up pathway, from ResNet
        enc0 = self.enc0(x)

        enc1 = self.enc1(enc0)  # 256

        enc2 = self.enc2(enc1)  # 512

        enc3 = self.enc3(enc2)  # 1024

        enc4 = self.enc4(enc3)  # 2048

        # Lateral connections

        lateral4 = self.lateral4(enc4)
        lateral3 = self.lateral3(enc3)
        lateral2 = self.lateral2(enc2)
        lateral1 = self.lateral1(enc1)
        lateral0 = self.lateral0(enc0)

        # Top-down pathway
        map4 = lateral4
        map3 = self.td1(
            lateral3 +
            nn.functional.upsample(map4, scale_factor=2, mode='nearest'))
        map2 = self.td2(
            lateral2 +
            nn.functional.upsample(map3, scale_factor=2, mode='nearest'))
        map1 = self.td3(
            lateral1 +
            nn.functional.upsample(map2, scale_factor=2, mode='nearest'))
        return lateral0, map1, map2, map3, map4
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor ): You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """
        # Bottom-up pathway, from ResNet
        enc0 = self.enc0(x)

        enc1 = self.enc1(enc0)  # 256

        enc2 = self.enc2(enc1)  # 512

        enc3 = self.enc3(enc2)  # 1024

        enc4 = self.enc4(enc3)  # 2048

        # Lateral connections

        lateral4 = self.lateral4(enc4)
        lateral3 = self.lateral3(enc3)
        lateral2 = self.lateral2(enc2)
        lateral1 = self.lateral1(enc1)
        lateral0 = self.lateral0(enc0)

        # Top-down pathway
        map4 = lateral4
        map3 = self.td1(
            lateral3 +
            nn.functional.interpolate(map4, scale_factor=2, mode='nearest'))
        map2 = self.td2(
            lateral2 +
            nn.functional.interpolate(map3, scale_factor=2, mode='nearest'))
        map1 = self.td3(
            lateral1 +
            nn.functional.interpolate(map2, scale_factor=2, mode='nearest'))
        return lateral0, map1, map2, map3, map4
//  [python]
def forward(self, x):
        content = x[0]
        condition = x[1]
        hw=(x[1].shape[2]//self.K)*(x[1].shape[3]//self.K)*x[1].shape[0]
        condition1=self.unfold(condition).view(x[1].shape[0],3,self.K,self.K,x[1].shape[2]//self.K,x[1].shape[3]//self.K)
        condition1=condition1.permute(0,4,5,1,2,3).view(hw,3,self.K,self.K)
        # conditiongt=condition[:,:,:3,:3]
        # conditionpatch0=condition1[:,:,:,:,0,0]
        # mse=torch.mean((conditiongt-conditionpatch0)*(conditiongt-conditionpatch0))
        # b,3,h,w
        fea = self.classifier(condition1).squeeze(2).squeeze(2)#[:,:,0,0].view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.cond_c)


        scale_first = self.cond_scale_first(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        shift_first = self.cond_shift_first(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        scale_first=F.upsample_bilinear(scale_first,(x[1].shape[2],x[1].shape[3]))
        shift_first=F.upsample_bilinear(shift_first,(x[1].shape[2],x[1].shape[3]))

        out = self.conv_first(content)
        out = out * scale_first + shift_first + out
        out = self.act(out)
        scale_first=scale_first.cpu()
        shift_first=shift_first.cpu()
        content=content.cpu()
        torch.cuda.empty_cache()


        scale_HR = self.cond_scale_HR(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        shift_HR = self.cond_shift_HR(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        scale_HR=F.upsample_bilinear(scale_HR,(x[1].shape[2],x[1].shape[3]))
        shift_HR=F.upsample_bilinear(shift_HR,(x[1].shape[2],x[1].shape[3]))
        
        out = self.HRconv(out)
        out = out * scale_HR + shift_HR + out
        out = self.act(out)
        
        scale_HR=scale_HR.cpu()
        shift_HR=shift_HR.cpu()
        torch.cuda.empty_cache()
        
        


        scale_last = self.cond_scale_last(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,3).permute(0,3,1,2)
        shift_last = self.cond_shift_last(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,3).permute(0,3,1,2)
        scale_last=F.upsample_bilinear(scale_last,(x[1].shape[2],x[1].shape[3]))
        shift_last=F.upsample_bilinear(shift_last,(x[1].shape[2],x[1].shape[3]))
        

        out = self.conv_last(out)
        out = out * scale_last + shift_last + out

        return out
//  [python]
def forward(self, x):
        content = x[0]
        condition = x[1]
        hw=(x[1].shape[2]//self.K)*(x[1].shape[3]//self.K)*x[1].shape[0]
        condition1=self.unfold(condition).view(x[1].shape[0],3,self.K,self.K,x[1].shape[2]//self.K,x[1].shape[3]//self.K)
        condition1=condition1.permute(0,4,5,1,2,3).view(hw,3,self.K,self.K)
        # conditiongt=condition[:,:,:3,:3]
        # conditionpatch0=condition1[:,:,:,:,0,0]
        # mse=torch.mean((conditiongt-conditionpatch0)*(conditiongt-conditionpatch0))
        # b,3,h,w
        fea = self.classifier(condition1).squeeze(2).squeeze(2)#[:,:,0,0].view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.cond_c)


        scale_first = self.cond_scale_first(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        shift_first = self.cond_shift_first(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        scale_first=F.interpolate(scale_first,(x[1].shape[2],x[1].shape[3]),mode='bilinear')
        shift_first=F.interpolate(shift_first,(x[1].shape[2],x[1].shape[3]),mode='bilinear')

        out = self.conv_first(content)
        out = out * scale_first + shift_first + out
        out = self.act(out)
        scale_first=scale_first.cpu()
        shift_first=shift_first.cpu()
        content=content.cpu()
        torch.cuda.empty_cache()


        scale_HR = self.cond_scale_HR(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        shift_HR = self.cond_shift_HR(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,self.nf).permute(0,3,1,2)
        scale_HR=F.interpolate(scale_HR,(x[1].shape[2],x[1].shape[3]),mode='bilinear')
        shift_HR=F.interpolate(shift_HR,(x[1].shape[2],x[1].shape[3]),mode='bilinear')
        
        out = self.HRconv(out)
        out = out * scale_HR + shift_HR + out
        out = self.act(out)
        
        scale_HR=scale_HR.cpu()
        shift_HR=shift_HR.cpu()
        torch.cuda.empty_cache()
        
        


        scale_last = self.cond_scale_last(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,3).permute(0,3,1,2)
        shift_last = self.cond_shift_last(fea).view(x[1].shape[0],x[1].shape[2]//self.K,x[1].shape[3]//self.K,3).permute(0,3,1,2)
        scale_last=F.interpolate(scale_last,(x[1].shape[2],x[1].shape[3]),mode='bilinear')
        shift_last=F.interpolate(shift_last,(x[1].shape[2],x[1].shape[3]),mode='bilinear')
        

        out = self.conv_last(out)
        out = out * scale_last + shift_last + out

        return out
//  [python]
def synth2real(self, noraml_img, noise=0.1, noise_patch=2, sigma=1.7, bk_value=0.95):
        device = noraml_img.device
        B, C, H, W = noraml_img.size()
        img_white_noise_scale = F.upsample_nearest(noise * torch.randn((B, C, H//noise_patch, W//noise_patch)).to(device), scale_factor=noise_patch) * noraml_img.lt(bk_value).float()
        noraml_img = noraml_img + img_white_noise_scale
        if sigma != 0:
            noraml_img = F.pad(noraml_img, (2, 2, 2, 2), mode='reflect')
            sigma = random.uniform(sigma, sigma)
            noraml_img = self.smoothing(noraml_img, sigma)
        return noraml_img
//  [python]
def synth2real(self, noraml_img, noise=0.1, noise_patch=2, sigma=1.7, bk_value=0.95):
        device = noraml_img.device
        B, C, H, W = noraml_img.size()
        img_white_noise_scale = F.interpolate(noise * torch.randn((B, C, H//noise_patch, W//noise_patch)).to(device), scale_factor=noise_patch, mode='nearest') * noraml_img.lt(bk_value).float()
        noraml_img = noraml_img + img_white_noise_scale
        if sigma != 0:
            noraml_img = F.pad(noraml_img, (2, 2, 2, 2), mode='reflect')
            sigma = random.uniform(sigma, sigma)
            noraml_img = self.smoothing(noraml_img, sigma)
        return noraml_img
//  [python]
def _upsample(self, x, scale=1):
        return F.upsample(x, scale_factor=scale, mode='bilinear')
//  [python]
def _upsample(self, x, scale=1):
        return F.interpolate(x, scale_factor=scale, mode='bilinear')
//  [python]
def _upsample_add(self, x, y, scale=1):
        return F.upsample(x, scale_factor=scale, mode='bilinear') + y
//  [python]
def _upsample_add(self, x, y, scale=1):
        return F.interpolate(x, scale_factor=scale, mode='bilinear') + y
//  [python]
def mask_from_positions(positions, box_size, ang_pix, distance):
    "generate mask from consensus gaussian model with mask being 1 if voxel center closer than distance (in angstrom) from a gaussian center"
    device = positions.device
    if box_size > 200:
        new_box = box_size//2
        new_ang_pix = ang_pix/2
        points = (positions*(box_size-1))/ang_pix
    else:
        new_box = box_size
        new_ang_pix = ang_pix
        points = (positions*(box_size-1))/ang_pix
    x = (torch.linspace(-0.5, 0.5, new_box)*(box_size-1))/ang_pix
    grid = torch.meshgrid(x, x, x, indexing='ij')
    grid = torch.stack([grid[0].ravel(), grid[1].ravel(), grid[2].ravel()], 1)
    tree = KDTree(points.detach().cpu().numpy())
    (dists, points) = tree.query(grid.cpu().numpy())
    ess_grid = grid[dists < distance]
    ess_grid_int = torch.clip(torch.round(
        (ess_grid - torch.min(grid))*new_ang_pix).long(), max=new_box-1)
    M = torch.zeros(new_box, new_box, new_box)
    M[ess_grid_int[:, 0], ess_grid_int[:, 1], ess_grid_int[:, 2]] = 1
    if box_size > 200:
        M = torch.nn.functional.upsample_nearest(M[None, None], scale_factor=2)
    return M[0, 0].to(device)
//  [python]
def mask_from_positions(positions, box_size, ang_pix, distance):
    "generate mask from consensus gaussian model with mask being 1 if voxel center closer than distance (in angstrom) from a gaussian center"
    device = positions.device
    if box_size > 200:
        new_box = box_size//2
        new_ang_pix = ang_pix/2
        points = (positions*(box_size-1))/ang_pix
    else:
        new_box = box_size
        new_ang_pix = ang_pix
        points = (positions*(box_size-1))/ang_pix
    x = (torch.linspace(-0.5, 0.5, new_box)*(box_size-1))/ang_pix
    grid = torch.meshgrid(x, x, x, indexing='ij')
    grid = torch.stack([grid[0].ravel(), grid[1].ravel(), grid[2].ravel()], 1)
    tree = KDTree(points.detach().cpu().numpy())
    (dists, points) = tree.query(grid.cpu().numpy())
    ess_grid = grid[dists < distance]
    ess_grid_int = torch.clip(torch.round(
        (ess_grid - torch.min(grid))*new_ang_pix).long(), max=new_box-1)
    M = torch.zeros(new_box, new_box, new_box)
    M[ess_grid_int[:, 0], ess_grid_int[:, 1], ess_grid_int[:, 2]] = 1
    if box_size > 200:
        M = torch.nn.functional.interpolate(M[None, None], scale_factor=2, mode='nearest')
    return M[0, 0].to(device)
//  [python]
def forward(self, input):
        '''
        :param input: Receives the input RGB image
        :return: the transformed feature map with spatial dimensions 1/8th of the input image
        '''
        output0 = self.level1(input)
        inp1 = self.sample1(input)
        inp2 = self.sample2(input)

        output0_cat = self.b1(torch.cat([output0, inp1], 1))
        output1_0 = self.level2_0(output0_cat) # down-sampled
        
        for i, layer in enumerate(self.level2):
            if i==0:
                output1 = layer(output1_0)
            else:
                output1 = layer(output1)

        output1_cat = self.b2(torch.cat([output1,  output1_0, inp2], 1))

        output2_0 = self.level3_0(output1_cat) # down-sampled
        for i, layer in enumerate(self.level3):
            if i==0:
                output2 = layer(output2_0)
            else:
                output2 = layer(output2)

        output2_cat = self.b3(torch.cat([output2_0, output2], 1))

        classifier = self.classifier(output2_cat)

        #return classifier
        out = F.upsample(classifier, input.size()[2:], mode='bilinear')   #Upsample score map, factor=8
        return out
//  [python]
def forward(self, input):
        '''
        :param input: Receives the input RGB image
        :return: the transformed feature map with spatial dimensions 1/8th of the input image
        '''
        output0 = self.level1(input)
        inp1 = self.sample1(input)
        inp2 = self.sample2(input)

        output0_cat = self.b1(torch.cat([output0, inp1], 1))
        output1_0 = self.level2_0(output0_cat) # down-sampled
        
        for i, layer in enumerate(self.level2):
            if i==0:
                output1 = layer(output1_0)
            else:
                output1 = layer(output1)

        output1_cat = self.b2(torch.cat([output1,  output1_0, inp2], 1))

        output2_0 = self.level3_0(output1_cat) # down-sampled
        for i, layer in enumerate(self.level3):
            if i==0:
                output2 = layer(output2_0)
            else:
                output2 = layer(output2)

        output2_cat = self.b3(torch.cat([output2_0, output2], 1))

        classifier = self.classifier(output2_cat)

        #return classifier
        out = F.interpolate(classifier, input.size()[2:], mode='bilinear')   #Upsample score map, factor=8
        return out
//  [python]
def _wanetTrigger(img, height, width, distance, trig_w, trig_h, delta=20, f=6):
    """
    Implement paper:
    > WaNet -- Imperceptible Warping-based Backdoor Attack
    > Anh Nguyen, Anh Tran, ICLR 2021
    > https://arxiv.org/abs/2102.10369
    """
    k = 4
    s = 0.5
    input_height = height
    grid_rescale = 1
    ins = torch.rand(1, 2, k, k) * 2 - 1
    ins = ins / torch.mean(torch.abs(ins))
    noise_grid = (F.upsample(ins,
                             size=input_height,
                             mode="bicubic",
                             align_corners=True).permute(0, 2, 3, 1))
    array1d = torch.linspace(-1, 1, steps=input_height)
    x, y = torch.meshgrid(array1d, array1d)
    # identity_grid = torch.stack((y, x), 2)[None, ...].to(device)
    identity_grid = torch.stack((y, x), 2)[None, ...]
    grid_temps = (identity_grid + s * noise_grid / input_height) * grid_rescale
    grid_temps = torch.clamp(grid_temps, -1, 1)
    img = np.float32(img)
    img = torch.tensor(img).reshape(-1, height, width).unsqueeze(0)
    img = F.grid_sample(img, grid_temps,
                        align_corners=True).squeeze(0).reshape(
                            height, width, -1)
    img = np.uint8(np.clip(img.cpu().numpy(), 0, 255))

    return img
//  [python]
def _wanetTrigger(img, height, width, distance, trig_w, trig_h, delta=20, f=6):
    """
    Implement paper:
    > WaNet -- Imperceptible Warping-based Backdoor Attack
    > Anh Nguyen, Anh Tran, ICLR 2021
    > https://arxiv.org/abs/2102.10369
    """
    k = 4
    s = 0.5
    input_height = height
    grid_rescale = 1
    ins = torch.rand(1, 2, k, k) * 2 - 1
    ins = ins / torch.mean(torch.abs(ins))
    noise_grid = (F.interpolate(ins,
                             size=input_height,
                             mode="bicubic",
                             align_corners=True).permute(0, 2, 3, 1))
    array1d = torch.linspace(-1, 1, steps=input_height)
    x, y = torch.meshgrid(array1d, array1d)
    # identity_grid = torch.stack((y, x), 2)[None, ...].to(device)
    identity_grid = torch.stack((y, x), 2)[None, ...]
    grid_temps = (identity_grid + s * noise_grid / input_height) * grid_rescale
    grid_temps = torch.clamp(grid_temps, -1, 1)
    img = np.float32(img)
    img = torch.tensor(img).reshape(-1, height, width).unsqueeze(0)
    img = F.grid_sample(img, grid_temps,
                        align_corners=True).squeeze(0).reshape(
                            height, width, -1)
    img = np.uint8(np.clip(img.cpu().numpy(), 0, 255))

    return img
//  [python]
def _wanetTriggerCross(img, height, width, distance, trig_w, trig_h):
    """
    Implement paper:
    > WaNet -- Imperceptible Warping-based Backdoor Attack
    > Anh Nguyen, Anh Tran, ICLR 2021
    > https://arxiv.org/abs/2102.10369
    """
    k = 4
    s = 0.5
    input_height = height
    grid_rescale = 1
    ins = torch.rand(1, 2, k, k) * 2 - 1
    ins = ins / torch.mean(torch.abs(ins))
    noise_grid = (F.upsample(ins,
                             size=input_height,
                             mode="bicubic",
                             align_corners=True).permute(0, 2, 3, 1))
    array1d = torch.linspace(-1, 1, steps=input_height)
    x, y = torch.meshgrid(array1d, array1d)
    identity_grid = torch.stack((y, x), 2)[None, ...]
    grid_temps = (identity_grid + s * noise_grid / input_height) * grid_rescale
    grid_temps = torch.clamp(grid_temps, -1, 1)
    ins = torch.rand(1, input_height, input_height, 2) * 2 - 1
    grid_temps2 = grid_temps + ins / input_height
    grid_temps2 = torch.clamp(grid_temps2, -1, 1)
    img = np.float32(img)
    img = torch.tensor(img).reshape(-1, height, width).unsqueeze(0)
    img = F.grid_sample(img, grid_temps2,
                        align_corners=True).squeeze(0).reshape(
                            height, width, -1)
    img = np.uint8(np.clip(img.cpu().numpy(), 0, 255))
    return img
//  [python]
def _wanetTriggerCross(img, height, width, distance, trig_w, trig_h):
    """
    Implement paper:
    > WaNet -- Imperceptible Warping-based Backdoor Attack
    > Anh Nguyen, Anh Tran, ICLR 2021
    > https://arxiv.org/abs/2102.10369
    """
    k = 4
    s = 0.5
    input_height = height
    grid_rescale = 1
    ins = torch.rand(1, 2, k, k) * 2 - 1
    ins = ins / torch.mean(torch.abs(ins))
    noise_grid = (F.interpolate(ins,
                             size=input_height,
                             mode="bicubic",
                             align_corners=True).permute(0, 2, 3, 1))
    array1d = torch.linspace(-1, 1, steps=input_height)
    x, y = torch.meshgrid(array1d, array1d)
    identity_grid = torch.stack((y, x), 2)[None, ...]
    grid_temps = (identity_grid + s * noise_grid / input_height) * grid_rescale
    grid_temps = torch.clamp(grid_temps, -1, 1)
    ins = torch.rand(1, input_height, input_height, 2) * 2 - 1
    grid_temps2 = grid_temps + ins / input_height
    grid_temps2 = torch.clamp(grid_temps2, -1, 1)
    img = np.float32(img)
    img = torch.tensor(img).reshape(-1, height, width).unsqueeze(0)
    img = F.grid_sample(img, grid_temps2,
                        align_corners=True).squeeze(0).reshape(
                            height, width, -1)
    img = np.uint8(np.clip(img.cpu().numpy(), 0, 255))
    return img
//  [python]
def forward(self, x, y):
        y = F.upsample(y, scale_factor=2)
        x = self.conv1(x)
        return self.conv2(self.padding2(x+y))
//  [python]
def forward(self, x, y):
        y = F.interpolate(y, scale_factor=2)
        x = self.conv1(x)
        return self.conv2(self.padding2(x+y))
//  [python]
def forward(self, x):
        x_size = x.size()
        x_conv1 = self.conv1(x)
        # print('x_conv1.shape:', x_conv1.shape)
        x = self.bn1(x_conv1)
        x = self.relu(x)
        x_pool = self.maxpool(x)

        x_layer1 = self.layer1(x_pool)
        # print('x_layer1.shape:', x_layer1.shape)
        x_layer2 = self.layer2(x_layer1)
        # print('x_layer2.shape:', x_layer2.shape)
        x = self.layer3(x_layer2)
        x = self.layer4(x)

        # ----for semantic segment----
        x_sem = self.out_conv(x)
        # print('x_sem.shape:', x_sem.shape)
        # x_sem = self.up(x_sem)
        x_sem = F.upsample_bilinear(x_sem, x_size[2:])
        # print('x_sem.shape:', x_sem.shape)
        # ----for semantic segment----

        # ----for object detection----
        x_det = self.layer5(x)
        x_det = self.conv_end(x_det)
        x_det = self.bn_end(x_det)
        x_det = F.sigmoid(x_det)
        x_det = x_det.permute(0,2,3,1)
        # ----for object detection----
        return x_sem, x_det
//  [python]
def forward(self, x):
        x_size = x.size()
        x_conv1 = self.conv1(x)
        # print('x_conv1.shape:', x_conv1.shape)
        x = self.bn1(x_conv1)
        x = self.relu(x)
        x_pool = self.maxpool(x)

        x_layer1 = self.layer1(x_pool)
        # print('x_layer1.shape:', x_layer1.shape)
        x_layer2 = self.layer2(x_layer1)
        # print('x_layer2.shape:', x_layer2.shape)
        x = self.layer3(x_layer2)
        x = self.layer4(x)

        # ----for semantic segment----
        x_sem = self.out_conv(x)
        # print('x_sem.shape:', x_sem.shape)
        # x_sem = self.up(x_sem)
        x_sem = F.interpolate(x_sem, x_size[2:], mode='bilinear')
        # print('x_sem.shape:', x_sem.shape)
        # ----for semantic segment----

        # ----for object detection----
        x_det = self.layer5(x)
        x_det = self.conv_end(x_det)
        x_det = self.bn_end(x_det)
        x_det = F.sigmoid(x_det)
        x_det = x_det.permute(0,2,3,1)
        # ----for object detection----
        return x_sem, x_det
//  [python]
def init_one_item(self,img_features, original_uv_images=None):
        # Find or init

        reduction_factor=4
        batch_size = img_features.shape[0]
        if self.get_len_dict() == 0:
            output_uv_imgs = 0.5*torch.randn((1,self.uv_c,self.uv_h//reduction_factor,self.uv_w//reduction_factor)).repeat(batch_size, 1, 1, 1)
            output_uv_imgs=torch.nn.functional.upsample_bilinear(output_uv_imgs, (self.uv_h,self.uv_w))
            #-original_uv_images
        else:
            for b in range(batch_size):
                img_feature=img_features
                dist = torch.sum(torch.pow(self.img_feature_dict-torch.unsqueeze(img_feature,0),2),1)
                min_idx = torch.argmin(dist,0)
                if b==0:
                    output_uv_imgs=torch.unsqueeze(self.uv_dict[min_idx],0)
                else:
                    output_uv_imgs=torch.cat((output_uv_imgs,torch.unsqueeze(self.uv_dict[min_idx],0)),0)
        return output_uv_imgs
//  [python]
def init_one_item(self,img_features, original_uv_images=None):
        # Find or init

        reduction_factor=4
        batch_size = img_features.shape[0]
        if self.get_len_dict() == 0:
            output_uv_imgs = 0.5*torch.randn((1,self.uv_c,self.uv_h//reduction_factor,self.uv_w//reduction_factor)).repeat(batch_size, 1, 1, 1)
            output_uv_imgs=torch.nn.functional.interpolate(output_uv_imgs, (self.uv_h,self.uv_w), mode='bilinear')
            #-original_uv_images
        else:
            for b in range(batch_size):
                img_feature=img_features
                dist = torch.sum(torch.pow(self.img_feature_dict-torch.unsqueeze(img_feature,0),2),1)
                min_idx = torch.argmin(dist,0)
                if b==0:
                    output_uv_imgs=torch.unsqueeze(self.uv_dict[min_idx],0)
                else:
                    output_uv_imgs=torch.cat((output_uv_imgs,torch.unsqueeze(self.uv_dict[min_idx],0)),0)
        return output_uv_imgs
//  [python]
def forward(self, X, label, requires_loss):
        """Forward pass of the network.
        """
        N = X.size()[0]
        X1 = self.features1(X)
        H = X1.size()[2]
        W = X1.size()[3]
        assert X1.size()[1] == 512
        X2 = self.features2(X)
        H2 = X2.size()[2]
        W2 = X2.size()[3]
        assert X2.size()[1] == 128

        if (H != H2) | (W != W2):
            X2 = F.upsample_bilinear(X2, (H, W))

        X1 = X1.view(N, 512, H * W)
        X2 = X2.view(N, 128, H * W)
        X = torch.bmm(X1, torch.transpose(X2, 1, 2)) / (H * W)  # Bilinear
        assert X.size() == (N, 512, 128)
        X = X.view(N, 512 * 128)
        X = torch.sqrt(X + 1e-8)
        X = torch.nn.functional.normalize(X)
        X = self.fc(X)
        assert X.size() == (N, 1)
        if requires_loss:
            return X, label, self.loss_build(X, label)
        else:
            return X