jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
                  cond_key=None, return_original_cond=False, bs=None):
        x = super().get_input(batch, k)
        if bs is not None:
            x = x[:bs]
        x = x.to(self.device)
        encoder_posterior = self.encode_first_stage(x)
        z = self.get_first_stage_encoding(encoder_posterior).detach()

        if self.model.conditioning_key is not None:
            if cond_key is None:
                cond_key = self.cond_stage_key
            if cond_key != self.first_stage_key:
                if cond_key in ['caption', 'coordinates_bbox']:
                    xc = batch[cond_key]
                elif cond_key == 'class_label':
                    xc = batch
                else:
                    xc = super().get_input(batch, cond_key).to(self.device)
            else:
                xc = x
            if not self.cond_stage_trainable or force_c_encode:
                if isinstance(xc, dict) or isinstance(xc, list):
                    # import pudb; pudb.set_trace()
                    c = self.get_learned_conditioning(xc)
                else:
                    c = self.get_learned_conditioning(xc.to(self.device))
            else:
                c = xc
            if bs is not None:
                c = c[:bs]
            # Testing #
            if cond_key == 'masked_image':
                mask = super().get_input(batch, "mask")
                cc = torch.nn.functional.interpolate(mask, size=c.shape[-2:]) # [B, 1, 10, 106]
                c = torch.cat((c, cc), dim=1) # [B, 5, 10, 106]
            # Testing #
            if self.use_positional_encodings:
                pos_x, pos_y = self.compute_latent_shifts(batch)
                ckey = __conditioning_keys__[self.model.conditioning_key]
                c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y}

        else:
            c = None
            xc = None
            if self.use_positional_encodings:
                pos_x, pos_y = self.compute_latent_shifts(batch)
                c = {'pos_x': pos_x, 'pos_y': pos_y}
        out = [z, c]
        if return_first_stage_outputs:
            xrec = self.decode_first_stage(z)
            out.extend([x, xrec])
        if return_original_cond:
            out.append(xc)
        return out
//  [python]
def __call__(self, face_img):
        H, W, C = face_img.shape

        w_size = 384
        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5

        x_data = x_data.to(self.cn_device)

        with torch.no_grad():
            hs = self.model(x_data[None, ...])
            heatmaps = torch.nn.functional.upsample(
                hs[-1],
                size=(H, W),
                mode='bilinear', align_corners=True).cpu().numpy()[0]
        return heatmaps
//  [python]
def __call__(self, face_img):
        H, W, C = face_img.shape

        w_size = 384
        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5

        x_data = x_data.to(self.cn_device)

        with torch.no_grad():
            hs = self.model(x_data[None, ...])
            heatmaps = F.interpolate(
                hs[-1],
                (H, W),
                mode='bilinear', align_corners=True).cpu().numpy()[0]
        return heatmaps
//  [python]
def forward(self, x):
        mask = F.max_pool2d(x, kernel_size=8, stride=4)
        mask = self.convs(mask)
        mask = F.upsample_bilinear(mask, size=(x.size(2), x.size(3)))
        mask = self.sigmoid(mask+x)
        
        return x * mask
//  [python]
def forward(self, x):
        mask = F.max_pool2d(x, kernel_size=8, stride=4)
        mask = self.convs(mask)
        mask = F.interpolate(mask, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) 
        mask = self.sigmoid(mask+x)
        
        return x * mask
//  [python]
def post_process_sem_seg_output(outputs: OneFormerForUniversalSegmentationOutput, target_size: Tuple[int, int]):
    # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1]
    class_queries_logits = outputs.class_queries_logits
    # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH]
    masks_queries_logits = outputs.masks_queries_logits
    if target_size is not None:
        masks_queries_logits = t.nn.functional.interpolate(
            masks_queries_logits,
            size=target_size,
            mode="bilinear",
            align_corners=False,
        )
    # remove the null class `[..., :-1]`
    masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
    # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH]
    masks_probs = masks_queries_logits.sigmoid()
    # now we want to sum over the queries,
    # $ out_{c,h,w} =  \sum_q p_{q,c} * m_{q,h,w} $
    # where $ softmax(p) \in R^{q, c} $ is the mask classes
    # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities
    # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth)
    segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)

    return segmentation
//  [python]
def post_process_sem_seg_output(outputs: OneFormerForUniversalSegmentationOutput, target_size: Tuple[int, int]):
    # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1]
    class_queries_logits = outputs.class_queries_logits
    # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH]
    masks_queries_logits = outputs.masks_queries_logits
    if target_size is not None:
        masks_queries_logits = torch.nn.functional.interpolate(
            masks_queries_logits,
            size=target_size,
            mode="bilinear",
            align_corners=False,
        )
    # remove the null class `[..., :-1]`
    masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
    # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH]
    masks_probs = masks_queries_logits.sigmoid()
    # now we want to sum over the queries,
    # $ out_{c,h,w} =  \sum_q p_{q,c} * m_{q,h,w} $
    # where $ softmax(p) \in R^{q, c} $ is the mask classes
    # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities
    # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth)
    segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)

    return segmentation
//  [python]
def __call__(self, imgs):
        img_augs = []
        for img in imgs:
            h, w = img.shape[:2]
            # later: provide list and randomly choose index for resize
            size = np.random.randint(self.short_edge_length[0], self.short_edge_length[1] + 1)
            if size == 0:
                return img
            scale = size * 1.0 / min(h, w)
            if h < w:
                newh, neww = size, scale * w
            else:
                newh, neww = scale * h, size
            if max(newh, neww) > self.max_size:
                scale = self.max_size * 1.0 / max(newh, neww)
                newh = newh * scale
                neww = neww * scale
            neww = int(neww + 0.5)
            newh = int(newh + 0.5)

            if img.dtype == np.uint8:
                pil_image = Image.fromarray(img)
                pil_image = pil_image.resize((neww, newh), Image.BILINEAR)
                img = np.asarray(pil_image)
            else:
                img = img.permute(2, 0, 1).unsqueeze(0)  # 3, 0, 1)  # hw(c) -> nchw
                img = t.interpolate(img, (newh, neww), mode=self.interp_method, align_corners=False).squeeze(0)
            img_augs.append(img)

        return img_augs
//  [python]
def __call__(self, imgs):
        img_augs = []
        for img in imgs:
            h, w = img.shape[:2]
            # later: provide list and randomly choose index for resize
            size = np.random.randint(self.short_edge_length[0], self.short_edge_length[1] + 1)
            if size == 0:
                return img
            scale = size * 1.0 / min(h, w)
            if h < w:
                newh, neww = size, scale * w
            else:
                newh, neww = scale * h, size
            if max(newh, neww) > self.max_size:
                scale = self.max_size * 1.0 / max(newh, neww)
                newh = newh * scale
                neww = neww * scale
            neww = int(neww + 0.5)
            newh = int(newh + 0.5)

            if img.dtype == np.uint8:
                pil_image = Image.fromarray(img)
                pil_image = pil_image.resize((neww, newh), Image.BILINEAR)
                img = np.asarray(pil_image)
            else:
                img = img.permute(2, 0, 1).unsqueeze(0)  # 3, 0, 1)  # hw(c) -> nchw
                img = F.interpolate(img, (newh, neww), mode=self.interp_method, align_corners=False).squeeze(0)
            img_augs.append(img)

        return img_augs
//  [python]
def forward(self, x):
        if x.size(2) != 256 or x.size(3) != 256:
            x = F.upsample(x, size=(256, 256), mode="bilinear")

        x = self.layer1(x)
        x = self.layer2(self.actvn(x))
        x = self.layer3(self.actvn(x))
        x = self.layer4(self.actvn(x))
        x = self.layer5(self.actvn(x))
        if self.opt.crop_size >= 256:
            x = self.layer6(self.actvn(x))
        x = self.actvn(x)

        x = x.view(x.size(0), -1)
        mu = self.fc_mu(x)
        logvar = self.fc_var(x)

        return mu, logvar
//  [python]
def forward(self, x):
        if x.size(2) != 256 or x.size(3) != 256:
            x = F.interpolate(x, size=(256, 256), mode="bilinear")

        x = self.layer1(x)
        x = self.layer2(self.actvn(x))
        x = self.layer3(self.actvn(x))
        x = self.layer4(self.actvn(x))
        x = self.layer5(self.actvn(x))
        if self.opt.crop_size >= 256:
            x = self.layer6(self.actvn(x))
        x = self.actvn(x)

        x = x.view(x.size(0), -1)
        mu = self.fc_mu(x)
        logvar = self.fc_var(x)

        return mu, logvar
//  [python]
def forward(self, x):
        s = self.act(self.squeeze(x))
        # Pyra 1
        d1 = self.act(self.down1(s))
        u1 = t.interpolate(d1, (s.size(2), s.size(3)), mode='bilinear', align_corners=False) 
        h1 = s - u1
        #h1 = u1

        # Pyra 2
        d2 = self.act(self.down2(d1))
        u2 = t.interpolate(d2, (d1.size(2), d1.size(3)), mode='bilinear', align_corners=False) 
        h2 = d1 - u2
        #h2 = u2
        h2 = t.interpolate(h2, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) 

        #'''
        # Pyra 3
        d3 = self.act(self.down3(d2))
        u3 = t.interpolate(d3, (d2.size(2), d2.size(3)), mode='bilinear', align_corners=False) 
        h3 = d2 - u3
        #h3 = u3
        h3 = t.interpolate(h3, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) 
        #'''

        m = self.excite(torch.cat((h1, h2, h3), dim=1))
        m = self.sigmoid(m)
        out = self.fuse(torch.cat((x*m, h1), dim=1))
        
        return out
//  [python]
def forward(self, x):
        s = self.act(self.squeeze(x))
        # Pyra 1
        d1 = self.act(self.down1(s))
        u1 = F.interpolate(d1, (s.size(2), s.size(3)), mode='bilinear', align_corners=False) 
        h1 = s - u1
        #h1 = u1

        # Pyra 2
        d2 = self.act(self.down2(d1))
        u2 = F.interpolate(d2, (d1.size(2), d1.size(3)), mode='bilinear', align_corners=False) 
        h2 = d1 - u2
        #h2 = u2
        h2 = F.interpolate(h2, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) 

        #'''
        # Pyra 3
        d3 = self.act(self.down3(d2))
        u3 = F.interpolate(d3, (d2.size(2), d2.size(3)), mode='bilinear', align_corners=False) 
        h3 = d2 - u3
        #h3 = u3
        h3 = F.interpolate(h3, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) 
        #'''

        m = self.excite(torch.cat((h1, h2, h3), dim=1))
        m = self.sigmoid(m)
        out = self.fuse(torch.cat((x*m, h1), dim=1))
        
        return out
//  [python]
def interpolate_groups(g, ratio, mode, align_corners):
    batch_size, num_objects = g.shape[:2]
    g = t.interpolate(g.flatten(start_dim=0, end_dim=1), 
                scale_factor=ratio, mode=mode, align_corners=align_corners)
    g = g.view(batch_size, num_objects, *g.shape[1:])
    return g
//  [python]
def interpolate_groups(g, ratio, mode, align_corners):
    batch_size, num_objects = g.shape[:2]
    g = F.interpolate(g.flatten(start_dim=0, end_dim=1), 
                scale_factor=ratio, mode=mode, align_corners=align_corners)
    g = g.view(batch_size, num_objects, *g.shape[1:])
    return g
//  [python]
def get_rel_pos(rel_pos, d):
    if isinstance(d, int):
        ori_d = rel_pos.shape[0]
        if ori_d == d:
            return rel_pos
        else:
            # Interpolate rel pos.
            new_pos_embed = t.nn.functional.interpolate(
                rel_pos.reshape(1, ori_d, -1).permute(0, 2, 1),
                size=d,
                mode="linear",
            )

            return new_pos_embed.reshape(-1, d).permute(1, 0)
//  [python]
def get_rel_pos(rel_pos, d):
    if isinstance(d, int):
        ori_d = rel_pos.shape[0]
        if ori_d == d:
            return rel_pos
        else:
            # Interpolate rel pos.
            new_pos_embed = F.interpolate(
                rel_pos.reshape(1, ori_d, -1).permute(0, 2, 1),
                size=d,
                mode="linear",
            )

            return new_pos_embed.reshape(-1, d).permute(1, 0)
//  [python]
def forward(self, input, input_lr):
        distilled_c1 = self.act(self.c1_d(input))
        r_c1 = (self.c1_r(input))
        r_c1 = self.act(r_c1+input)

        distilled_c2 = self.act(self.c2_d(r_c1))
        r_c2 = (self.c2_r(r_c1))
        r_c2 = self.act(r_c2+r_c1)

        distilled_c3 = self.act(self.c3_d(r_c2))
        r_c3 = (self.c3_r(r_c2))
        r_c3 = self.act(r_c3+r_c2)

        r_c4 = self.act(self.c4(r_c3))

        out = torch.cat([distilled_c1, distilled_c2, distilled_c3, r_c4], dim=1)
        out_fused = (self.c5(out)) 

        lr_feat = self.lr_convs(input_lr)
        hr_feat = F.max_pool2d(out_fused, kernel_size=4, stride=4)
        mask = torch.cat([lr_feat, hr_feat], dim=1)
        mask = self.mask_convs(mask)
        mask = F.upsample(mask, size=(input.size(2), input.size(3)), mode='bilinear') 
        out_masked = out_fused * mask


        return out_masked, lr_feat
//  [python]
def forward(self, input, input_lr):
        distilled_c1 = self.act(self.c1_d(input))
        r_c1 = (self.c1_r(input))
        r_c1 = self.act(r_c1+input)

        distilled_c2 = self.act(self.c2_d(r_c1))
        r_c2 = (self.c2_r(r_c1))
        r_c2 = self.act(r_c2+r_c1)

        distilled_c3 = self.act(self.c3_d(r_c2))
        r_c3 = (self.c3_r(r_c2))
        r_c3 = self.act(r_c3+r_c2)

        r_c4 = self.act(self.c4(r_c3))

        out = torch.cat([distilled_c1, distilled_c2, distilled_c3, r_c4], dim=1)
        out_fused = (self.c5(out)) 

        lr_feat = self.lr_convs(input_lr)
        hr_feat = F.max_pool2d(out_fused, kernel_size=4, stride=4)
        mask = torch.cat([lr_feat, hr_feat], dim=1)
        mask = self.mask_convs(mask)
        mask = F.interpolate(mask, (input.size(2), input.size(3)), mode='bilinear', align_corners=False) 
        out_masked = out_fused * mask


        return out_masked, lr_feat
//  [python]
def test_upsamplingLinear1d(self):
        for align_corners in [True, False]:
            for recompute_scale_factor in [True, False]:
                kwargs = dict(
                    mode='linear', align_corners=align_corners, recompute_scale_factor=recompute_scale_factor
                )
                # test float scale factor up & downsampling
                for scale_factor in [0.5, 1.5, 2]:
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs)
                    in_t = torch.ones(1, 1, 2)
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    with warnings.catch_warnings(record=True) as w:
                        out_t = m(in_t)
                    self.assertEqual(torch.ones(1, 1, out_size), out_t.data)

                    input1 = torch.randn(1, 1, 2, requires_grad=True)
                    if not recompute_scale_factor:
                        gradcheck(lambda x: t.interpolate(x, out_size, **kwargs), (input1,))
                    else:
                        gradcheck(lambda x: t.interpolate(x, scale_factor=scale_factor, **kwargs), (input1,))
//  [python]
def test_upsamplingLinear1d(self):
        for align_corners in [True, False]:
            for recompute_scale_factor in [True, False]:
                kwargs = dict(
                    mode='linear', align_corners=align_corners, recompute_scale_factor=recompute_scale_factor
                )
                # test float scale factor up & downsampling
                for scale_factor in [0.5, 1.5, 2]:
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs)
                    in_t = torch.ones(1, 1, 2)
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    with warnings.catch_warnings(record=True) as w:
                        out_t = m(in_t)
                    self.assertEqual(torch.ones(1, 1, out_size), out_t.data)

                    input1 = torch.randn(1, 1, 2, requires_grad=True)
                    if not recompute_scale_factor:
                        gradcheck(lambda x: F.interpolate(x, out_size, **kwargs), (input1,))
                    else:
                        gradcheck(lambda x: F.interpolate(x, scale_factor=scale_factor, **kwargs), (input1,))
//  [python]
def test_interpolate_illegal_memory_access(self):
        in_s = 45
        out_s = 14

        input1 = torch.ones((1, 1, in_s), device='npu', requires_grad=True)
        # note we allocated grad_output to be larger so out of bound access
        # would be visible in grad_input
        grad = torch.ones((1, 1, out_s * 2), device='npu', requires_grad=True)
        grad = grad[:, :, :out_s]

        input_ref = input1.detach().cpu().requires_grad_()
        grad_ref = grad.cpu()

        out = t.nn.functional.interpolate(input1, size=(out_s,), mode='nearest')
        out.backward(grad)

        out_ref = t.nn.functional.interpolate(input_ref, size=(out_s,), mode='nearest')
        out_ref.backward(grad_ref)

        self.assertEqual(out_ref, out)
        self.assertEqual(input_ref.grad, input1.grad)
//  [python]
def test_interpolate_illegal_memory_access(self):
        in_s = 45
        out_s = 14

        input1 = torch.ones((1, 1, in_s), device='npu', requires_grad=True)
        # note we allocated grad_output to be larger so out of bound access
        # would be visible in grad_input
        grad = torch.ones((1, 1, out_s * 2), device='npu', requires_grad=True)
        grad = grad[:, :, :out_s]

        input_ref = input1.detach().cpu().requires_grad_()
        grad_ref = grad.cpu()

        out = F.interpolate(input1, size=(out_s,), mode='nearest')
        out.backward(grad)

        out_ref = F.interpolate(input_ref, size=(out_s,), mode='nearest')
        out_ref.backward(grad_ref)

        self.assertEqual(out_ref, out)
        self.assertEqual(input_ref.grad, input1.grad)
//  [python]
def test_upsamplingNearestExact1d_correctness(self, device, isize, osize):
        in_t = torch.arange(isize, dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
        out_t = t.nn.functional.interpolate(
            in_t, size=(osize, ), recompute_scale_factor=False, mode="nearest-exact"
        )
        # compute expected output as scikit-image/scipy
        expected_out = torch.zeros(osize, dtype=torch.float).unsqueeze(0).unsqueeze(0)
        scale = 1.0 * isize / osize
        for m in range(osize):
            i_f32 = (m + 0.5) * scale
            i = int(i_f32)
            expected_out[0, 0, m] = in_t[0, 0, i]
        expected_out = expected_out.to(device=device)
        self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact1d_correctness(self, device, isize, osize):
        in_t = torch.arange(isize, dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
        out_t = F.interpolate(
            in_t, size=(osize, ), recompute_scale_factor=False, mode="nearest-exact"
        )
        # compute expected output as scikit-image/scipy
        expected_out = torch.zeros(osize, dtype=torch.float).unsqueeze(0).unsqueeze(0)
        scale = 1.0 * isize / osize
        for m in range(osize):
            i_f32 = (m + 0.5) * scale
            i = int(i_f32)
            expected_out[0, 0, m] = in_t[0, 0, i]
        expected_out = expected_out.to(device=device)
        self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact1d_rescale(self, device):
        isize = 20
        in_t = torch.arange(isize, dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
        for s in [1.00001, ]:
            out_t = t.nn.functional.interpolate(
                in_t, scale_factor=s, recompute_scale_factor=False, mode="nearest-exact"
            )
            expected_out = in_t
            self.assertEqual(out_t, expected_out, msg=f"scale: {s}")

        for s in [2.00001, ]:
            out_t = t.nn.functional.interpolate(
                in_t, scale_factor=s, recompute_scale_factor=False, mode="nearest-exact"
            )
            # input is [[[0, 1, 2, 3, ..., 9]]]
            # expected out is [[[0, 0, 1, 1, 2, 2, ..., 9, 9]]]
            expected_out = in_t.repeat_interleave(2, dim=-1)
            self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact1d_rescale(self, device):
        isize = 20
        in_t = torch.arange(isize, dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
        for s in [1.00001, ]:
            out_t = F.interpolate(
                in_t, scale_factor=s, recompute_scale_factor=False, mode="nearest-exact"
            )
            expected_out = in_t
            self.assertEqual(out_t, expected_out, msg=f"scale: {s}")

        for s in [2.00001, ]:
            out_t = F.interpolate(
                in_t, scale_factor=s, recompute_scale_factor=False, mode="nearest-exact"
            )
            # input is [[[0, 1, 2, 3, ..., 9]]]
            # expected out is [[[0, 0, 1, 1, 2, 2, ..., 9, 9]]]
            expected_out = in_t.repeat_interleave(2, dim=-1)
            self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact2d_correctness(self, device, memory_format, isize, osize):
        in_t = torch.arange(isize * isize, dtype=torch.float, device=device).reshape(1, 1, isize, isize)
        in_t = in_t.contiguous(memory_format=memory_format)
        out_t = t.nn.functional.interpolate(
            in_t, size=(osize, osize), recompute_scale_factor=False, mode="nearest-exact"
        )
        # compute expected output as Scikit-Image/Scipy
        expected_out = torch.zeros(1, 1, osize, osize, dtype=torch.float)
        scale = 1.0 * isize / osize
        for o1 in range(osize):
            i1_f32 = (o1 + 0.5) * scale
            i1 = int(i1_f32)
            for o2 in range(osize):
                i2_f32 = (o2 + 0.5) * scale
                i2 = int(i2_f32)
                expected_out[0, 0, o1, o2] = in_t[0, 0, i1, i2]
        expected_out = expected_out.to(device=device)
        self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact2d_correctness(self, device, memory_format, isize, osize):
        in_t = torch.arange(isize * isize, dtype=torch.float, device=device).reshape(1, 1, isize, isize)
        in_t = in_t.contiguous(memory_format=memory_format)
        out_t = F.interpolate(
            in_t, size=(osize, osize), recompute_scale_factor=False, mode="nearest-exact"
        )
        # compute expected output as Scikit-Image/Scipy
        expected_out = torch.zeros(1, 1, osize, osize, dtype=torch.float)
        scale = 1.0 * isize / osize
        for o1 in range(osize):
            i1_f32 = (o1 + 0.5) * scale
            i1 = int(i1_f32)
            for o2 in range(osize):
                i2_f32 = (o2 + 0.5) * scale
                i2 = int(i2_f32)
                expected_out[0, 0, o1, o2] = in_t[0, 0, i1, i2]
        expected_out = expected_out.to(device=device)
        self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact3d_correctness(self, device, memory_format, isize, osize):
        in_t = torch.arange(isize * isize * isize, dtype=torch.float, device=device)
        in_t = in_t.reshape(1, 1, isize, isize, isize)
        in_t = in_t.contiguous(memory_format=memory_format)
        out_t = t.nn.functional.interpolate(
            in_t, size=(osize, osize, osize), recompute_scale_factor=False, mode="nearest-exact"
        )
        # compute expected output as Scikit-Image/Scipy
        expected_out = torch.zeros(1, 1, osize, osize, osize, dtype=torch.float)
        scale = 1.0 * isize / osize
        for o1 in range(osize):
            i1_f32 = (o1 + 0.5) * scale
            i1 = int(i1_f32)
            for o2 in range(osize):
                i2_f32 = (o2 + 0.5) * scale
                i2 = int(i2_f32)
                for o3 in range(osize):
                    i3_f32 = (o3 + 0.5) * scale
                    i3 = int(i3_f32)
                    expected_out[0, 0, o1, o2, o3] = in_t[0, 0, i1, i2, i3]
        expected_out = expected_out.to(device=device)
        self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingNearestExact3d_correctness(self, device, memory_format, isize, osize):
        in_t = torch.arange(isize * isize * isize, dtype=torch.float, device=device)
        in_t = in_t.reshape(1, 1, isize, isize, isize)
        in_t = in_t.contiguous(memory_format=memory_format)
        out_t = F.interpolate(
            in_t, size=(osize, osize, osize), recompute_scale_factor=False, mode="nearest-exact"
        )
        # compute expected output as Scikit-Image/Scipy
        expected_out = torch.zeros(1, 1, osize, osize, osize, dtype=torch.float)
        scale = 1.0 * isize / osize
        for o1 in range(osize):
            i1_f32 = (o1 + 0.5) * scale
            i1 = int(i1_f32)
            for o2 in range(osize):
                i2_f32 = (o2 + 0.5) * scale
                i2 = int(i2_f32)
                for o3 in range(osize):
                    i3_f32 = (o3 + 0.5) * scale
                    i3 = int(i3_f32)
                    expected_out[0, 0, o1, o2, o3] = in_t[0, 0, i1, i2, i3]
        expected_out = expected_out.to(device=device)
        self.assertEqual(out_t, expected_out)
//  [python]
def test_upsamplingBiLinear2d_consistency_interp_size_bug(self, device, memory_format, align_corners, input_size, output_size):

        if torch.device(device).type == "npu":
            raise SkipTest("NPU implementation is not yet supporting uint8")

        mode = "bilinear"
        input_ui8 = torch.randint(0, 256, size=(1, 3, input_size, input_size), dtype=torch.uint8, device=device)
        input_ui8 = input_ui8.contiguous(memory_format=memory_format)
        input_f32 = input_ui8.float()

        output_f32 = t.interpolate(
            input_f32, size=(output_size, output_size), mode=mode, align_corners=align_corners
        ).round().to(torch.uint8)
        output_ui8 = t.interpolate(
            input_ui8, size=(output_size, output_size), mode=mode, align_corners=align_corners
        )
        torch.testing.assert_close(output_f32, output_ui8, atol=1, rtol=0)
//  [python]
def test_upsamplingBiLinear2d_consistency_interp_size_bug(self, device, memory_format, align_corners, input_size, output_size):

        if torch.device(device).type == "npu":
            raise SkipTest("NPU implementation is not yet supporting uint8")

        mode = "bilinear"
        input_ui8 = torch.randint(0, 256, size=(1, 3, input_size, input_size), dtype=torch.uint8, device=device)
        input_ui8 = input_ui8.contiguous(memory_format=memory_format)
        input_f32 = input_ui8.float()

        output_f32 = F.interpolate(
            input_f32, size=(output_size, output_size), mode=mode, align_corners=align_corners, antialias=False
        ).round().to(torch.uint8)
        output_ui8 = F.interpolate(
            input_ui8, size=(output_size, output_size), mode=mode, align_corners=align_corners, antialias=False
        )
        torch.testing.assert_close(output_f32, output_ui8, atol=1, rtol=0)
//  [python]
def get_view(self,view_index):
        idx=-1
        for list_idx, idx_value in enumerate(self.view_index_list):
            if abs(idx_value - view_index.item()) < 1e-5:
                idx = list_idx
                break
        assert idx!=-1
        idx += random.randint(0,self.per_view_num-1)

        cur_path = self.image_path[idx]
        view_index = self.view_index_list[idx]
        cur_image = PIL.Image.open(cur_path)
        image = np.array(cur_image, dtype=np.float32)  # (h,w,c)
        image /= 255.0  #
        image = np.transpose(image, (2, 0, 1))  # (c, h, w)
        image = torch.from_numpy(image)  # numpy->tensor
        image = t.interpolate(image[None], (512, 512), mode='bilinear', align_corners=False)[0]
        return image
//  [python]
def get_view(self,view_index):
        idx=-1
        for list_idx, idx_value in enumerate(self.view_index_list):
            if abs(idx_value - view_index.item()) < 1e-5:
                idx = list_idx
                break
        assert idx!=-1
        idx += random.randint(0,self.per_view_num-1)

        cur_path = self.image_path[idx]
        view_index = self.view_index_list[idx]
        cur_image = PIL.Image.open(cur_path)
        image = np.array(cur_image, dtype=np.float32)  # (h,w,c)
        image /= 255.0  #
        image = np.transpose(image, (2, 0, 1))  # (c, h, w)
        image = torch.from_numpy(image)  # numpy->tensor
        image = F.interpolate(image[None], (512, 512), mode='bilinear', align_corners=False)[0]
        return image
//  [python]
def __getitem__(self, idx):
        cur_path = self.image_path[idx]
        view_index = self.view_index_list[idx]
        cur_image = PIL.Image.open(cur_path)
        image = np.array(cur_image, dtype=np.float32)  # (h,w,c)
        image /= 255.0  #
        image = np.transpose(image, (2, 0, 1))  # (c, h, w)
        image = torch.from_numpy(image)  # numpy->tensor
        image = t.interpolate(image[None], (512, 512), mode='bilinear', align_corners=False)[0]

        return image, view_index
//  [python]
def __getitem__(self, idx):
        cur_path = self.image_path[idx]
        view_index = self.view_index_list[idx]
        cur_image = PIL.Image.open(cur_path)
        image = np.array(cur_image, dtype=np.float32)  # (h,w,c)
        image /= 255.0  #
        image = np.transpose(image, (2, 0, 1))  # (c, h, w)
        image = torch.from_numpy(image)  # numpy->tensor
        image = F.interpolate(image[None], (512, 512), mode='bilinear', align_corners=False)[0]

        return image, view_index
//  [python]
def test_interpolate_buffer_overflow(self):

        def helper(size, dtype, mode, device, is_channels_last):
            input1 = torch.ones(size, dtype=dtype, device=device)
            if is_channels_last:
                if len(size) == 3:
                    input1 = input1.transpose(1, 2).contiguous().transpose(1, 2)
                elif len(size) == 4:
                    input1 = input1.to(memory_format=torch.channels_last)
                else:
                    input1 = input1.to(memory_format=torch.channels_last_3d)
            output1 = t.interpolate(input1, 2, mode=mode, align_corners=True)
            # reset the corner value and expect the output is changed as well
            # the output won't be changed on buffer overflow
            input1[(-1,) * len(size)] = 0.5
            output2 = t.interpolate(input1, 2, mode=mode, align_corners=True)
            self.assertNotEqual(output1, output2)

        size_dtype_list = []
        # We set the size larger than the floating point exactly representable range
        # float: exact representable range (-2**24,2**24)
        size_dtype_list.append(([1, 10, 2**24 + 4], torch.float))
        size_dtype_list.append(([1, 10, 2, 2**24 + 4], torch.float))
        size_dtype_list.append(([1, 10, 2, 2, 2**24 + 4], torch.float))
        # bfloat16: exact representable range (-2**8, 2**8)
        size_dtype_list.append(([1, 10, 2**8 + 4], torch.bfloat16))
        size_dtype_list.append(([1, 10, 2, 2**8 + 4], torch.bfloat16))
        size_dtype_list.append(([1, 10, 2, 2, 2**8 + 4], torch.bfloat16))
        # half: exact representable range (-2**11, 2**11)
        size_dtype_list.append(([1, 10, 2**11 + 4], torch.half))
        size_dtype_list.append(([1, 10, 2, 2**11 + 4], torch.half))
        size_dtype_list.append(([1, 10, 2, 2, 2**11 + 4], torch.half))

        # devices = ['cpu'] + (['cuda'] if torch.cuda.is_available() else [])
        devices = ['cpu']

        for mode in ('linear', 'bilinear', 'bicubic', 'trilinear'):
            for size_dtype in size_dtype_list:
                size, dtype = size_dtype
                if (
                    mode == 'linear' and len(size) != 3
                    or (mode == 'bilinear' and len(size) != 4)
                    or (mode == 'bicubic' and len(size) != 4)
                    or (mode == 'trilinear' and len(size) != 5)
                ):
                    continue
                for device in devices:
                    if (
                        device == 'cpu' and dtype == torch.half
                        or (device == 'npu' and dtype == torch.bfloat16)
                    ):
                        # no half precision support on cpu or bfloat16 on npu yet
                        continue
                    for is_channels_last in (True, False):
                        helper(size, dtype, mode, device, is_channels_last)
//  [python]
def test_interpolate_buffer_overflow(self):

        def helper(size, dtype, mode, device, is_channels_last):
            input1 = torch.ones(size, dtype=dtype, device=device)
            if is_channels_last:
                if len(size) == 3:
                    input1 = input1.transpose(1, 2).contiguous().transpose(1, 2)
                elif len(size) == 4:
                    input1 = input1.to(memory_format=torch.channels_last)
                else:
                    input1 = input1.to(memory_format=torch.channels_last_3d)
            output1 = F.interpolate(input1, 2, mode=mode, align_corners=True)
            # reset the corner value and expect the output is changed as well
            # the output won't be changed on buffer overflow
            input1[(-1,) * len(size)] = 0.5
            output2 = F.interpolate(input1, 2, mode=mode, align_corners=True)
            self.assertNotEqual(output1, output2)

        size_dtype_list = []
        # We set the size larger than the floating point exactly representable range
        # float: exact representable range (-2**24,2**24)
        size_dtype_list.append(([1, 10, 2**24 + 4], torch.float))
        size_dtype_list.append(([1, 10, 2, 2**24 + 4], torch.float))
        size_dtype_list.append(([1, 10, 2, 2, 2**24 + 4], torch.float))
        # bfloat16: exact representable range (-2**8, 2**8)
        size_dtype_list.append(([1, 10, 2**8 + 4], torch.bfloat16))
        size_dtype_list.append(([1, 10, 2, 2**8 + 4], torch.bfloat16))
        size_dtype_list.append(([1, 10, 2, 2, 2**8 + 4], torch.bfloat16))
        # half: exact representable range (-2**11, 2**11)
        size_dtype_list.append(([1, 10, 2**11 + 4], torch.half))
        size_dtype_list.append(([1, 10, 2, 2**11 + 4], torch.half))
        size_dtype_list.append(([1, 10, 2, 2, 2**11 + 4], torch.half))

        # devices = ['cpu'] + (['cuda'] if torch.cuda.is_available() else [])
        devices = ['cpu']

        for mode in ('linear', 'bilinear', 'bicubic', 'trilinear'):
            for size_dtype in size_dtype_list:
                size, dtype = size_dtype
                if (
                    mode == 'linear' and len(size) != 3
                    or (mode == 'bilinear' and len(size) != 4)
                    or (mode == 'bicubic' and len(size) != 4)
                    or (mode == 'trilinear' and len(size) != 5)
                ):
                    continue
                for device in devices:
                    if (
                        device == 'cpu' and dtype == torch.half
                        or (device == 'npu' and dtype == torch.bfloat16)
                    ):
                        # no half precision support on cpu or bfloat16 on npu yet
                        continue
                    for is_channels_last in (True, False):
                        helper(size, dtype, mode, device, is_channels_last)
//  [python]
def test_upsamplingTrilinear3d(self, device, align_corners):
        kwargs = dict(mode='trilinear', align_corners=align_corners)

        for memory_format in [torch.contiguous_format, torch.channels_last_3d]:
            # test float scale factor up & downsampling
            for scale_factor in [0.5, 1.5, 2]:
                m = nn.Upsample(scale_factor=scale_factor, **kwargs)
                in_t = torch.ones(1, 2, 2, 2, 2, device=device, dtype=torch.double)
                in_t = in_t.contiguous(memory_format=memory_format).requires_grad_()
                out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                with warnings.catch_warnings(record=True) as w:
                    out_t = m(in_t)
                expected_out = torch.ones(1, 2, out_size, out_size, out_size, device=device, dtype=torch.double)
                self.assertEqual(expected_out, out_t)
                # Assert that memory format is carried through to the output
                self.assertTrue(out_t.is_contiguous(memory_format=memory_format))
                out_t.backward(torch.randn_like(out_t))
                self.assertTrue(in_t.grad.is_contiguous(memory_format=memory_format))

                input1 = torch.randn(1, 2, 2, 2, 2, requires_grad=True, dtype=torch.double)
                self.assertEqual(
                    t.interpolate(input1, (out_size, out_size, out_size), **kwargs),
                    t.interpolate(input1, scale_factor=scale_factor, **kwargs))
                gradcheck(lambda x: t.interpolate(x, out_size, **kwargs), [input1])
                gradgradcheck(lambda x: t.interpolate(x, out_size, **kwargs), [input1])
//  [python]
def test_upsamplingTrilinear3d(self, device, align_corners):
        kwargs = dict(mode='trilinear', align_corners=align_corners)

        for memory_format in [torch.contiguous_format, torch.channels_last_3d]:
            # test float scale factor up & downsampling
            for scale_factor in [0.5, 1.5, 2]:
                m = nn.Upsample(scale_factor=scale_factor, **kwargs)
                in_t = torch.ones(1, 2, 2, 2, 2, device=device, dtype=torch.double)
                in_t = in_t.contiguous(memory_format=memory_format).requires_grad_()
                out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                with warnings.catch_warnings(record=True) as w:
                    out_t = m(in_t)
                expected_out = torch.ones(1, 2, out_size, out_size, out_size, device=device, dtype=torch.double)
                self.assertEqual(expected_out, out_t)
                # Assert that memory format is carried through to the output
                self.assertTrue(out_t.is_contiguous(memory_format=memory_format))
                out_t.backward(torch.randn_like(out_t))
                self.assertTrue(in_t.grad.is_contiguous(memory_format=memory_format))

                input1 = torch.randn(1, 2, 2, 2, 2, requires_grad=True, dtype=torch.double)
                self.assertEqual(
                    F.interpolate(input1, (out_size, out_size, out_size), **kwargs),
                    F.interpolate(input1, scale_factor=scale_factor, **kwargs))
                gradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [input1])
                gradgradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [input1])
//  [python]
def upsample_grid(self, resolution: int) -> None:
        """Upsamples underyling feature grid

        Args:
            resolution: Target resolution.
        """

        line_coef = t.upsample_bilinear(self.line_coef.data, size=(resolution, 1))
        self.line_coef = torch.nn.Parameter(line_coef)

        self.resolution = resolution
//  [python]
def upsample_grid(self, resolution: int) -> None:
        """Upsamples underyling feature grid

        Args:
            resolution: Target resolution.
        """

        line_coef = F.interpolate(self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True)
        self.line_coef = torch.nn.Parameter(line_coef)

        self.resolution = resolution
//  [python]
def upsample_grid(self, resolution: int) -> None:
        """Upsamples underlying feature grid

        Args:
            resolution: Target resolution.
        """
        plane_coef = t.interpolate(
            self.plane_coef.data, size=(resolution, resolution), mode="bilinear", align_corners=True
        )
        line_coef = t.interpolate(self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True)

        self.plane_coef, self.line_coef = torch.nn.Parameter(plane_coef), torch.nn.Parameter(line_coef)
        self.resolution = resolution
//  [python]
def upsample_grid(self, resolution: int) -> None:
        """Upsamples underlying feature grid

        Args:
            resolution: Target resolution.
        """
        plane_coef = F.interpolate(
            self.plane_coef.data, size=(resolution, resolution), mode="bilinear", align_corners=True
        )
        line_coef = F.interpolate(self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True)

        self.plane_coef, self.line_coef = torch.nn.Parameter(plane_coef), torch.nn.Parameter(line_coef)
        self.resolution = resolution
//  [python]
def l2_loss(self):

        try:
            fake_data = next(self.d_iter)
        except Exception as e:
            self.d_iter = iter(self.d_loader)
            fake_data = next(self.d_iter)
        rays_o = fake_data['rays_o']  # [B, N, 3]
        rays_d = fake_data['rays_d']  # [B, N, 3]
        mvp = fake_data['mvp']  # [B, 4, 4]
        view_index = fake_data['view_index']

        B, N = rays_o.shape[:2]
        H, W = fake_data['H'], fake_data['W']
        ambient_ratio = 1.0
        shading = 'lambertian'

        rand = random.random()
        if (self.opt.bg_radius > 0 and rand > 0.5) or self.opt.no_random_bkg or self.opt.g_bkg_network:
            bg_color = None  # use bg_net
        else:
            bg_color = torch.rand(3).to(self.device)  # single color random bg

        outputs = self.model.render(rays_o, rays_d, mvp, H, W, staged = False, perturb = True, bg_color=bg_color, ambient_ratio=ambient_ratio, shading=shading)
        loss = 0
        loss_dict = {}

        fake_rgb = outputs['image'].reshape(B, H, W, 3).permute(0, 3, 1, 2).contiguous()  # [B, 3, H, W]
        fake_rgb = t.nn.functional.interpolate(fake_rgb, (512, 512), mode='bilinear', align_corners=False)

        real_data = self.real_dataset.get_view(view_index)
        real_data = real_data.to(self.device)

        l2_loss = self.mse(fake_rgb, real_data).mean()

        loss_dict['l2_loss'] = l2_loss.item()
        loss += l2_loss * self.cur_l2_weight
        return loss, loss_dict
//  [python]
def l2_loss(self):

        try:
            fake_data = next(self.d_iter)
        except Exception as e:
            self.d_iter = iter(self.d_loader)
            fake_data = next(self.d_iter)
        rays_o = fake_data['rays_o']  # [B, N, 3]
        rays_d = fake_data['rays_d']  # [B, N, 3]
        mvp = fake_data['mvp']  # [B, 4, 4]
        view_index = fake_data['view_index']

        B, N = rays_o.shape[:2]
        H, W = fake_data['H'], fake_data['W']
        ambient_ratio = 1.0
        shading = 'lambertian'

        rand = random.random()
        if (self.opt.bg_radius > 0 and rand > 0.5) or self.opt.no_random_bkg or self.opt.g_bkg_network:
            bg_color = None  # use bg_net
        else:
            bg_color = torch.rand(3).to(self.device)  # single color random bg

        outputs = self.model.render(rays_o, rays_d, mvp, H, W, staged = False, perturb = True, bg_color=bg_color, ambient_ratio=ambient_ratio, shading=shading)
        loss = 0
        loss_dict = {}

        fake_rgb = outputs['image'].reshape(B, H, W, 3).permute(0, 3, 1, 2).contiguous()  # [B, 3, H, W]
        fake_rgb = F.interpolate(fake_rgb, (512, 512), mode='bilinear', align_corners=False)

        real_data = self.real_dataset.get_view(view_index)
        real_data = real_data.to(self.device)

        l2_loss = self.mse(fake_rgb, real_data).mean()

        loss_dict['l2_loss'] = l2_loss.item()
        loss += l2_loss * self.cur_l2_weight
        return loss, loss_dict
//  [python]
def d_main_loss(self, fake_data, real_image, real_view_index):
        rays_o = fake_data['rays_o'] # [B, N, 3]
        rays_d = fake_data['rays_d'] # [B, N, 3]
        mvp = fake_data['mvp'] # [B, 4, 4]
        view_index = fake_data['view_index']
        view_input = torch.tensor([view_index,0.0],device=self.device)
        B, N = rays_o.shape[:2]
        H, W = fake_data['H'], fake_data['W']
        ambient_ratio = 1.0
        # render a
        shading = 'lambertian'

        binarize = False
        bg_color = None
        with torch.no_grad():
            outputs = self.model.render(rays_o, rays_d, mvp, H, W, staged = True, perturb = False, bg_color=bg_color, ambient_ratio=ambient_ratio, shading=shading, binarize=binarize)
            # perturb false related to Nerf, should be false when train D, True when train G
        fake_rgb = outputs['image'].reshape(B, H, W, 3).permute(0, 3, 1, 2).contiguous()  # [B, 3, H, W]
        fake_rgb = F.upsample(fake_rgb, size=(512, 512), mode='bilinear')
        fake_logits = self.query_D(fake_rgb, view_input)
        loss_fake = (torch.nn.functional.softplus(fake_logits)).mean()

        # show_image(outputs['image'].reshape(B, H, W, 3))
        real_image = real_image.to(self.device)
        real_view_input = torch.tensor([real_view_index,0.0],device=self.device)
        real_logits = self.query_D(real_image, real_view_input)
        loss_real = (torch.nn.functional.softplus(-real_logits)).mean()

        return loss_fake, loss_real
//  [python]
def d_main_loss(self, fake_data, real_image, real_view_index):
        rays_o = fake_data['rays_o'] # [B, N, 3]
        rays_d = fake_data['rays_d'] # [B, N, 3]
        mvp = fake_data['mvp'] # [B, 4, 4]
        view_index = fake_data['view_index']
        view_input = torch.tensor([view_index,0.0],device=self.device)
        B, N = rays_o.shape[:2]
        H, W = fake_data['H'], fake_data['W']
        ambient_ratio = 1.0
        # render a
        shading = 'lambertian'

        binarize = False
        bg_color = None
        with torch.no_grad():
            outputs = self.model.render(rays_o, rays_d, mvp, H, W, staged = True, perturb = False, bg_color=bg_color, ambient_ratio=ambient_ratio, shading=shading, binarize=binarize)
            # perturb false related to Nerf, should be false when train D, True when train G
        fake_rgb = outputs['image'].reshape(B, H, W, 3).permute(0, 3, 1, 2).contiguous()  # [B, 3, H, W]
        fake_rgb = F.interpolate(fake_rgb, (512, 512), mode='bilinear', align_corners=False)
        fake_logits = self.query_D(fake_rgb, view_input)
        loss_fake = (torch.nn.functional.softplus(fake_logits)).mean()

        # show_image(outputs['image'].reshape(B, H, W, 3))
        real_image = real_image.to(self.device)
        real_view_input = torch.tensor([real_view_index,0.0],device=self.device)
        real_logits = self.query_D(real_image, real_view_input)
        loss_real = (torch.nn.functional.softplus(-real_logits)).mean()

        return loss_fake, loss_real
//  [python]
def repeat_expand(
        self,
        content: Union[torch.Tensor, np.ndarray],
        target_len: int,
        mode: str = "nearest",
    ):
        ndim = content.ndim

        if content.ndim == 1:
            content = content[None, None]
        elif content.ndim == 2:
            content = content[None]

        assert content.ndim == 3

        is_np = isinstance(content, np.ndarray)
        if is_np:
            content = torch.from_numpy(content)

        results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)

        if is_np:
            results = results.numpy()

        if ndim == 1:
            return results[0, 0]
        elif ndim == 2:
            return results[0]