jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def retouch_local(self, image):
        """
        image: rgb
        """
        with torch.no_grad():
            sub_H, sub_W = image.shape[2:]

            sub_image_standard = F.interpolate(
                image, size=(768, 768), mode='bilinear', align_corners=True)
            sub_mask_pred = torch.sigmoid(
                self.detection_net(sub_image_standard))
            sub_mask_pred = F.interpolate(
                sub_mask_pred, size=(sub_H, sub_W), mode='nearest')

            sub_mask_pred_hard_low = (sub_mask_pred >= 0.35).float()
            sub_mask_pred_hard_high = (sub_mask_pred >= 0.5).float()
            sub_mask_pred = sub_mask_pred * (
                1 - sub_mask_pred_hard_high) + sub_mask_pred_hard_high
            sub_mask_pred = sub_mask_pred * sub_mask_pred_hard_low
            sub_mask_pred = 1 - sub_mask_pred

            sub_H_standard = sub_H if sub_H % self.patch_size == 0 else (
                sub_H // self.patch_size + 1) * self.patch_size
            sub_W_standard = sub_W if sub_W % self.patch_size == 0 else (
                sub_W // self.patch_size + 1) * self.patch_size

            sub_image_padding = F.pad(
                image,
                pad=(0, sub_W_standard - sub_W, 0, sub_H_standard - sub_H, 0,
                     0),
                mode='constant',
                value=0)
            sub_mask_pred_padding = F.pad(
                sub_mask_pred,
                pad=(0, sub_W_standard - sub_W, 0, sub_H_standard - sub_H, 0,
                     0),
                mode='constant',
                value=0)

            sub_image_padding = patch_partition_overlap(
                sub_image_padding, p1=self.patch_size, p2=self.patch_size)
            sub_mask_pred_padding = patch_partition_overlap(
                sub_mask_pred_padding, p1=self.patch_size, p2=self.patch_size)
            B_padding, C_padding, _, _ = sub_image_padding.size()

            sub_comp_padding_list = []
            for window_item in range(B_padding):
                sub_image_padding_window = sub_image_padding[
                    window_item:window_item + 1]
                sub_mask_pred_padding_window = sub_mask_pred_padding[
                    window_item:window_item + 1]

                sub_input_image_padding_window = sub_image_padding_window * sub_mask_pred_padding_window

                sub_output_padding_window = self.inpainting_net(
                    sub_input_image_padding_window,
                    sub_mask_pred_padding_window)
                sub_comp_padding_window = sub_input_image_padding_window + (
                    1
                    - sub_mask_pred_padding_window) * sub_output_padding_window

                sub_comp_padding_list.append(sub_comp_padding_window)

            sub_comp_padding = torch.cat(sub_comp_padding_list, dim=0)
            sub_comp = patch_aggregation_overlap(
                sub_comp_padding,
                h=int(round(sub_H_standard / self.patch_size)),
                w=int(round(sub_W_standard
                            / self.patch_size)))[:, :, :sub_H, :sub_W]

            return sub_comp
//  [python]
def vuv_binary(self, f0):
        """Calculate V/UV binary sequences.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: V/UV binary sequences (B, 1, T).

        """
        _, _, T = f0.size()
        uv = t.nn.functional.interpolate((f0 > 0) * torch.ones_like(f0), size=T * self.hop_size, mode='nearest')
        return uv
//  [python]
def vuv_binary(self, f0):
        """Calculate V/UV binary sequences.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: V/UV binary sequences (B, 1, T).

        """
        _, _, T = f0.size()
        uv = interpolate((f0 > 0) * torch.ones_like(f0), T * self.hop_size)
        return uv
//  [python]
def test_group_norm_backward(self, device='mps'):
        # See https://github.com/pytorch/pytorch/issues/88331 for more detail
        shape = [1, 4, 16, 16]
        x = torch.full(shape, 7.0, device=device)

        target = torch.ones((1, 3, 128, 128), device=device)

        conv_in = nn.Conv2d(4, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), device=device)
        conv_out = nn.Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), device=device)
        norm = nn.GroupNorm(32, 128, eps=1e-6, affine=True, device=device)

        with torch.enable_grad():
            x = x.detach().requires_grad_()
            out = 5.5 * x
            out = conv_in(out)
            out = out + norm(out)
            out = out + norm(out)
            out = out + norm(out)
            out = F.interpolate(out, scale_factor=8.0, mode="nearest")
            out = norm(out)
            out = conv_out(out)

            loss = (out - target).norm(dim=-1).sum()
            grad = -torch.autograd.grad(loss, x)[0]
            self.assertFalse(grad.detach().isnan().any().item(), 'NaN gradients returned by autograd')
//  [python]
def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, C, F, T).

        Returns:
            Tensor: Interpolated tensor (B, C, F * y_scale, T * x_scale)

        """
        return F.interpolate(x, size=(int(x.shape[2] * self.y_scale), int(x.shape[3] * self.x_scale)), mode=self.mode)
//  [python]
def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Input tensor (B, C, F, T).

        Returns:
            Tensor: Interpolated tensor (B, C, F * y_scale, T * x_scale)

        """
        return F.interpolate(x, scale_factor=(self.y_scale, self.x_scale), mode=self.mode)
//  [python]
def sinusoid(self, f0):
        """Calculate sine signals.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: Sines generated following NSF (B, 1, T).

        """
        B, _, T = f0.size()
        vuv = t.nn.functional.interpolate((f0 > 0) * torch.ones_like(f0), T * self.hop_size)
        radious = t.nn.functional.interpolate(f0, T * self.hop_size) / self.sample_rate % 1
        sine = vuv * torch.sin(torch.cumsum(radious, dim=2) * 2 * np.pi) * self.sine_amp
        if self.noise_amp > 0:
            noise_amp = vuv * self.noise_amp + (1.0 - vuv) * self.noise_amp / 3.0
            noise = torch.randn((B, 1, T * self.hop_size), device=f0.device) * noise_amp
            sine = sine + noise
        return sine
//  [python]
def sinusoid(self, f0):
        """Calculate sine signals.

        Args:
            f0 (Tensor): F0 tensor (B, 1, T // hop_size).

        Returns:
            Tensor: Sines generated following NSF (B, 1, T).

        """
        B, _, T = f0.size()
        vuv = interpolate((f0 > 0) * torch.ones_like(f0), T * self.hop_size)
        radious = interpolate(f0, T * self.hop_size) / self.sample_rate % 1
        sine = vuv * torch.sin(torch.cumsum(radious, dim=2) * 2 * np.pi) * self.sine_amp
        if self.noise_amp > 0:
            noise_amp = vuv * self.noise_amp + (1.0 - vuv) * self.noise_amp / 3.0
            noise = torch.randn((B, 1, T * self.hop_size), device=f0.device) * noise_amp
            sine = sine + noise
        return sine
//  [python]
def forward(self, x: torch.Tensor):
        """
        :param x: is the input feature map with shape `[batch_size, channels, height, width]`
        """
        # Up-sample by a factor of $2$
        x = F.upsample(x, scale_factor=2, mode="nearest")
        # Apply convolution
        return self.conv(x)
//  [python]
def forward(self, x: torch.Tensor):
        """
        :param x: is the input feature map with shape `[batch_size, channels, height, width]`
        """
        # Up-sample by a factor of $2$
        x = F.interpolate(x, scale_factor=2, mode="nearest")
        # Apply convolution
        return self.conv(x)
//  [python]
def forward(self, xt, xl):
        _, _, h, w = xl.size()

        xt = F.upsample(self.conv1(xt), size=(h, w), mode='bilinear')
        xl = self.conv2(xl)
        x = torch.cat([xt, xl], dim=1)
        x = self.conv3(x)
        seg = self.conv4(x)
        return seg, x
//  [python]
def forward(self, xt, xl):
        _, _, h, w = xl.size()

        xt = F.interpolate(self.conv1(xt), size=(h, w), mode='bilinear', align_corners=True)
        xl = self.conv2(xl)
        x = torch.cat([xt, xl], dim=1)
        x = self.conv3(x)
        seg = self.conv4(x)
        return seg, x
//  [python]
def test_interpolate(self):
        def helper(shape, output_size, scales, mode, align_corners=False):
            inputCPU = torch.randn(shape, device='cpu', dtype=torch.float, requires_grad=True)
            inputCPU.retain_grad()
            inputMPS = inputCPU.detach().clone().to('mps').requires_grad_()

            # align_corners is used for 2D interpolation only
            if (align_corners is True and len(shape) > 3 and mode == 'bilinear'):
                if scales is not None:
                    outputCPU = t.interpolate(inputCPU, scale_factor=scales, mode=mode, align_corners=align_corners)
                    outputMPS = t.interpolate(inputMPS, scale_factor=scales, mode=mode, align_corners=align_corners)
                else:
                    outputCPU = t.interpolate(inputCPU, size=output_size, mode=mode, align_corners=align_corners)
                    outputMPS = t.interpolate(inputMPS, size=output_size, mode=mode, align_corners=align_corners)
            elif scales is not None:
                outputCPU = t.interpolate(inputCPU, scale_factor=scales, mode=mode)
                outputMPS = t.interpolate(inputMPS, scale_factor=scales, mode=mode)
            else:
                outputCPU = t.interpolate(inputCPU, size=output_size, mode=mode)
                outputMPS = t.interpolate(inputMPS, size=output_size, mode=mode)

            self.assertEqual(outputCPU, outputMPS)

            # backward pass (chose 0.6 just to have the grad_output != 1)
            outputCPU.backward(gradient=torch.full_like(outputCPU, 0.6))
            outputMPS.backward(gradient=torch.full_like(outputMPS, 0.6))
            self.assertEqual(inputCPU.grad, inputMPS.grad)

        # 1D interpolation
        for mode in ['nearest', 'nearest-exact']:
            helper([2, 3, 4], [3], None, mode)  # downsample with size
            helper([2, 3, 4], [6], None, mode)  # upsample with size
            helper([2, 3, 4], None, [0.6], mode)  # downsample with scale factor
            helper([2, 3, 4], None, [1.7], mode)  # upsample with scale factor
        # 2D interpolation
        for mode in ['nearest', 'nearest-exact', 'bilinear']:
            helper([2, 3, 4, 5], [3, 4], None, mode)  # downsample_nearest with size
            helper([2, 3, 4, 5], [6, 7], None, mode)  # upsample_nearest with size
            helper([2, 3, 4, 5], None, [0.6, 0.7], mode)  # downsample_nearest with scale factor
            helper([2, 3, 4, 5], None, [1.4, 1.7], mode)  # upsample_nearest with scale factor
        # align_corners=True
        helper([2, 3, 4, 5], [3, 4], None, 'bilinear', True)
        helper([2, 3, 4, 5], None, [1.4, 1.7], 'bilinear', True)
//  [python]
def test_interpolate(self):
        def helper(shape, output_size, scales, mode, align_corners=False):
            inputCPU = torch.randn(shape, device='cpu', dtype=torch.float, requires_grad=True)
            inputCPU.retain_grad()
            inputMPS = inputCPU.detach().clone().to('mps').requires_grad_()

            # align_corners is used for 2D interpolation only
            if (align_corners is True and len(shape) > 3 and mode == 'bilinear'):
                if scales is not None:
                    outputCPU = nn.functional.interpolate(inputCPU, scale_factor=scales, mode=mode, align_corners=align_corners)
                    outputMPS = nn.functional.interpolate(inputMPS, scale_factor=scales, mode=mode, align_corners=align_corners)
                else:
                    outputCPU = nn.functional.interpolate(inputCPU, size=output_size, mode=mode, align_corners=align_corners)
                    outputMPS = nn.functional.interpolate(inputMPS, size=output_size, mode=mode, align_corners=align_corners)
            elif scales is not None:
                outputCPU = nn.functional.interpolate(inputCPU, scale_factor=scales, mode=mode)
                outputMPS = nn.functional.interpolate(inputMPS, scale_factor=scales, mode=mode)
            else:
                outputCPU = nn.functional.interpolate(inputCPU, size=output_size, mode=mode)
                outputMPS = nn.functional.interpolate(inputMPS, size=output_size, mode=mode)

            self.assertEqual(outputCPU, outputMPS)

            # backward pass (chose 0.6 just to have the grad_output != 1)
            outputCPU.backward(gradient=torch.full_like(outputCPU, 0.6))
            outputMPS.backward(gradient=torch.full_like(outputMPS, 0.6))
            self.assertEqual(inputCPU.grad, inputMPS.grad)

        # 1D interpolation
        for mode in ['nearest', 'nearest-exact']:
            helper([2, 3, 4], [3], None, mode)  # downsample with size
            helper([2, 3, 4], [6], None, mode)  # upsample with size
            helper([2, 3, 4], None, [0.6], mode)  # downsample with scale factor
            helper([2, 3, 4], None, [1.7], mode)  # upsample with scale factor
        # 2D interpolation
        for mode in ['nearest', 'nearest-exact', 'bilinear']:
            helper([2, 3, 4, 5], [3, 4], None, mode)  # downsample_nearest with size
            helper([2, 3, 4, 5], [6, 7], None, mode)  # upsample_nearest with size
            helper([2, 3, 4, 5], None, [0.6, 0.7], mode)  # downsample_nearest with scale factor
            helper([2, 3, 4, 5], None, [1.4, 1.7], mode)  # upsample_nearest with scale factor
        # align_corners=True
        helper([2, 3, 4, 5], [3, 4], None, 'bilinear', True)
        helper([2, 3, 4, 5], None, [1.4, 1.7], 'bilinear', True)
//  [python]
def get_batch(latents, list_ind, offset):
        '''prepare a batch out of the list of latents'''
        samples = latents[list_ind]['samples']
        shape = samples.shape
        mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu')
        if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]:
            t.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")
        if mask.shape[0] < samples.shape[0]:
            mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]
        if 'batch_index' in latents[list_ind]:
            batch_inds = latents[list_ind]['batch_index']
        else:
            batch_inds = [x+offset for x in range(shape[0])]
        return samples, mask, batch_inds
//  [python]
def get_batch(latents, list_ind, offset):
        '''prepare a batch out of the list of latents'''
        samples = latents[list_ind]['samples']
        shape = samples.shape
        mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu')
        if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]:
            torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")
        if mask.shape[0] < samples.shape[0]:
            mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]
        if 'batch_index' in latents[list_ind]:
            batch_inds = latents[list_ind]['batch_index']
        else:
            batch_inds = [x+offset for x in range(shape[0])]
        return samples, mask, batch_inds
//  [python]
def forward(self, x):
    if self.num_branches == 1:
      return [self.branches[0](x[0])]

    for i in range(self.num_branches):
      x[i] = self.branches[i](x[i])

    x_fuse = []
    for i in range(len(self.fuse_layers)):
      y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
      for j in range(1, self.num_branches):
        if i == j:
          y = y + x[j]
        elif j > i:
          y = y + t(self.fuse_layers[i][j](x[j]), size=[x[i].shape[2], x[i].shape[3]], mode='bilinear', align_corners=False)
        else:
          y = y + self.fuse_layers[i][j](x[j])
      x_fuse.append(self.relu(y))

    return x_fuse
//  [python]
def forward(self, x):
    if self.num_branches == 1:
      return [self.branches[0](x[0])]

    for i in range(self.num_branches):
      x[i] = self.branches[i](x[i])

    x_fuse = []
    for i in range(len(self.fuse_layers)):
      y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
      for j in range(1, self.num_branches):
        if i == j:
          y = y + x[j]
        elif j > i:
          y = y + F.interpolate(self.fuse_layers[i][j](x[j]), size=[x[i].shape[2], x[i].shape[3]], mode='bilinear', align_corners=False)
        else:
          y = y + self.fuse_layers[i][j](x[j])
      x_fuse.append(self.relu(y))

    return x_fuse
//  [python]
def adaptive_instance_normalization(content_feat, style_feat, c_mask=None, s_mask=None):
    assert (content_feat.size()[:2] == style_feat.size()[:2])

    size = content_feat.size()
    H, W = size[2], size[3]
    msk = None
    if c_mask is not None:
        msk = t.interpolate(c_mask, (H, W))

    s_size = style_feat.size()
    s_H, s_W = s_size[2], s_size[3]
    s_msk = None
    if s_mask is not None:
        s_msk = t.interpolate(s_mask, (s_H, s_W))

    style_mean, style_std = calc_mean_std(style_feat, mask=s_msk)
    content_mean, content_std = calc_mean_std(content_feat, mask=msk)

    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)

    if msk is not None:
        return (normalized_feat * style_std.expand(size) + style_mean.expand(size)) * msk + content_feat * (1 - msk)
    else:
        return normalized_feat * style_std.expand(size) + style_mean.expand(size)
//  [python]
def adaptive_instance_normalization(content_feat, style_feat, c_mask=None, s_mask=None):
    assert (content_feat.size()[:2] == style_feat.size()[:2])

    size = content_feat.size()
    H, W = size[2], size[3]
    msk = None
    if c_mask is not None:
        msk = F.interpolate(c_mask, (H, W))

    s_size = style_feat.size()
    s_H, s_W = s_size[2], s_size[3]
    s_msk = None
    if s_mask is not None:
        s_msk = F.interpolate(s_mask, (s_H, s_W))

    style_mean, style_std = calc_mean_std(style_feat, mask=s_msk)
    content_mean, content_std = calc_mean_std(content_feat, mask=msk)

    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)

    if msk is not None:
        return (normalized_feat * style_std.expand(size) + style_mean.expand(size)) * msk + content_feat * (1 - msk)
    else:
        return normalized_feat * style_std.expand(size) + style_mean.expand(size)
//  [python]
def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
    assert os.path.isfile(checkpoint_path)
    checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")

    saved_state_dict = checkpoint_dict["model"]
    if hasattr(model, "module"):
        state_dict = model.module.state_dict()
    else:
        state_dict = model.state_dict()
    new_state_dict = {}
    for k, v in state_dict.items():  # 模型需要的shape
        try:
            new_state_dict[k] = saved_state_dict[k]
            if saved_state_dict[k].shape != state_dict[k].shape:
                print(
                    f"shape-{k}-mismatch|need-{state_dict[k].shape}|get-{saved_state_dict[k].shape}"
                )
                if saved_state_dict[k].dim() == 2:  # NOTE: check is this ok?
                    # for embedded input 256 <==> 768
                    # this achieves we can continue training from original's pretrained checkpoints when using embedder that 768-th dim output etc.
                    if saved_state_dict[k].dtype == torch.half:
                        new_state_dict[k] = (
                            t.interpolate(
                                saved_state_dict[k].float().unsqueeze(0).unsqueeze(0),
                                size=state_dict[k].shape,
                                mode="bilinear",
                            )
                            .half()
                            .squeeze(0)
                            .squeeze(0)
                        )
                    else:
                        new_state_dict[k] = (
                            t.interpolate(
                                saved_state_dict[k].unsqueeze(0).unsqueeze(0),
                                size=state_dict[k].shape,
                                mode="bilinear",
                            )
                            .squeeze(0)
                            .squeeze(0)
                        )
                    print(
                        "interpolated new_state_dict",
                        k,
                        "from",
                        saved_state_dict[k].shape,
                        "to",
                        new_state_dict[k].shape,
                    )
                else:
                    raise KeyError
        except Exception as e:
            # print(traceback.format_exc())
            print(f"{k} is not in the checkpoint")
            print("error: %s" % e)
            new_state_dict[k] = v  # 模型自带的随机值
    if hasattr(model, "module"):
        model.module.load_state_dict(new_state_dict, strict=False)
    else:
        model.load_state_dict(new_state_dict, strict=False)
    print("Loaded model weights")

    epoch = checkpoint_dict["epoch"]
    learning_rate = checkpoint_dict["learning_rate"]
    if optimizer is not None and load_opt == 1:
        optimizer.load_state_dict(checkpoint_dict["optimizer"])
    print("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, epoch))
    return model, optimizer, learning_rate, epoch
//  [python]
def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
    assert os.path.isfile(checkpoint_path)
    checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")

    saved_state_dict = checkpoint_dict["model"]
    if hasattr(model, "module"):
        state_dict = model.module.state_dict()
    else:
        state_dict = model.state_dict()
    new_state_dict = {}
    for k, v in state_dict.items():  # 模型需要的shape
        try:
            new_state_dict[k] = saved_state_dict[k]
            if saved_state_dict[k].shape != state_dict[k].shape:
                print(
                    f"shape-{k}-mismatch|need-{state_dict[k].shape}|get-{saved_state_dict[k].shape}"
                )
                if saved_state_dict[k].dim() == 2:  # NOTE: check is this ok?
                    # for embedded input 256 <==> 768
                    # this achieves we can continue training from original's pretrained checkpoints when using embedder that 768-th dim output etc.
                    if saved_state_dict[k].dtype == torch.half:
                        new_state_dict[k] = (
                            F.interpolate(
                                saved_state_dict[k].float().unsqueeze(0).unsqueeze(0),
                                size=state_dict[k].shape,
                                mode="bilinear",
                            )
                            .half()
                            .squeeze(0)
                            .squeeze(0)
                        )
                    else:
                        new_state_dict[k] = (
                            F.interpolate(
                                saved_state_dict[k].unsqueeze(0).unsqueeze(0),
                                size=state_dict[k].shape,
                                mode="bilinear",
                            )
                            .squeeze(0)
                            .squeeze(0)
                        )
                    print(
                        "interpolated new_state_dict",
                        k,
                        "from",
                        saved_state_dict[k].shape,
                        "to",
                        new_state_dict[k].shape,
                    )
                else:
                    raise KeyError
        except Exception as e:
            # print(traceback.format_exc())
            print(f"{k} is not in the checkpoint")
            print("error: %s" % e)
            new_state_dict[k] = v  # 模型自带的随机值
    if hasattr(model, "module"):
        model.module.load_state_dict(new_state_dict, strict=False)
    else:
        model.load_state_dict(new_state_dict, strict=False)
    print("Loaded model weights")

    epoch = checkpoint_dict["epoch"]
    learning_rate = checkpoint_dict["learning_rate"]
    if optimizer is not None and load_opt == 1:
        optimizer.load_state_dict(checkpoint_dict["optimizer"])
    print("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, epoch))
    return model, optimizer, learning_rate, epoch
//  [python]
def postprocess(self, model_outputs):
        predicted_depth = model_outputs.predicted_depth
        prediction = torch.nn.functional.upsample(
            predicted_depth.unsqueeze(1), size=self.image_size[::-1], mode="bicubic", align_corners=False
        )
        output = prediction.squeeze().cpu().numpy()
        formatted = (output * 255 / np.max(output)).astype("uint8")
        depth = Image.fromarray(formatted)
        output_dict = {}
        output_dict["predicted_depth"] = predicted_depth
        output_dict["depth"] = depth
        return output_dict
//  [python]
def postprocess(self, model_outputs):
        predicted_depth = model_outputs.predicted_depth
        prediction = torch.nn.functional.interpolate(
            predicted_depth.unsqueeze(1), size=self.image_size[::-1], mode="bicubic", align_corners=False
        )
        output = prediction.squeeze().cpu().numpy()
        formatted = (output * 255 / np.max(output)).astype("uint8")
        depth = Image.fromarray(formatted)
        output_dict = {}
        output_dict["predicted_depth"] = predicted_depth
        output_dict["depth"] = depth
        return output_dict
//  [python]
def _interpolate(batch_image: torch.Tensor,
                 ds_factor: int,
                 mode: str = 'bicubic'
                 ) -> Tuple[torch.Tensor, Tuple[int, int]]:
    """Resize image to multiple of give down-sampling factor."""

    h, w = batch_image.shape[-2:]  # NCHW

    new_h = h - (h % ds_factor)
    new_w = w - (w % ds_factor)

    size = (new_h, new_w)
    if new_h != h or new_w != w:
        batch_image = t.interpolate(batch_image, size=size, mode=mode)

    return batch_image, size
//  [python]
def _interpolate(batch_image: torch.Tensor,
                 ds_factor: int,
                 mode: str = 'bicubic'
                 ) -> Tuple[torch.Tensor, Tuple[int, int]]:
    """Resize image to multiple of give down-sampling factor."""

    h, w = batch_image.shape[-2:]  # NCHW

    new_h = h - (h % ds_factor)
    new_w = w - (w % ds_factor)

    size = (new_h, new_w)
    if new_h != h or new_w != w:
        batch_image = F.interpolate(batch_image, size=size, mode=mode)

    return batch_image, size
//  [python]
def forward(self, m):
        if self.use_aux_net:
            aux = self.resnet(m)
            aux = t.interpolate(aux, scale_factor=self.scale, mode="linear", align_corners=True)
            aux = aux.transpose(1, 2)
        else:
            aux = None
        m = t.interpolate(m, scale_factor=self.scale, mode="linear", align_corners=True)
        m = m[:, :, self.indent : -self.indent]
        m = m * 0.045  # empirically found

        return m.transpose(1, 2), aux
//  [python]
def forward(self, m):
        if self.use_aux_net:
            aux = self.resnet(m)
            aux = torch.nn.functional.interpolate(aux, scale_factor=self.scale, mode="linear", align_corners=True)
            aux = aux.transpose(1, 2)
        else:
            aux = None
        m = torch.nn.functional.interpolate(m, scale_factor=self.scale, mode="linear", align_corners=True)
        m = m[:, :, self.indent : -self.indent]
        m = m * 0.045  # empirically found

        return m.transpose(1, 2), aux
//  [python]
def preprocess_image(self, batched_inputs, training):
        """
        Normalize, pad and batch the input images.
        """
        images = [x["image"].to(self.device) for x in batched_inputs]
        bs = len(images)
        images = [self.normalizer(x) for x in images]

        images = ImageList.from_tensors(
            images, size_divisibility=0, pad_ref_long=True)

        # sync image size for all gpus
        comm.synchronize()
        if training and self.iter % self.change_iter == 0:
            if self.iter < self.max_iter - 20000:
                meg = torch.LongTensor(1).to(self.device)
                comm.synchronize()
                if comm.is_main_process():
                    size = np.random.choice(self.multi_size)
                    meg.fill_(size)

                if comm.get_world_size() > 1:
                    comm.synchronize()
                    dist.broadcast(meg, 0)
                self.size = meg.item()

                comm.synchronize()
            else:
                self.size = 608

        if training:

            # resize image inputs
            modes = ['bilinear', 'nearest', 'bicubic', 'area']
            mode = modes[random.randrange(4)]
            images.tensor = t.nn.functional.interpolate(
                    images.tensor, size=[self.size, self.size], mode=mode, align_corners=False)

            if "instances" in batched_inputs[0]:
                gt_instances = [
                    x["instances"].to(self.device) for x in batched_inputs
                ]
            elif "targets" in batched_inputs[0]:
                log_first_n(
                    logging.WARN,
                    "'targets' in the model inputs is now renamed to 'instances'!",
                    n=10)
                gt_instances = [
                    x["targets"].to(self.device) for x in batched_inputs
                ]
            else:
                gt_instances = None

            targets = [
                torch.cat(
                    [instance.gt_classes.float().unsqueeze(-1), instance.gt_boxes.tensor], dim=-1
                )
                for instance in gt_instances
            ]
            labels = torch.zeros((bs, 100, 5))
            for i, target in enumerate(targets):
                labels[i][:target.shape[0]] = target
            labels[:, :, 1:] = labels[:, :, 1:] / 512. * self.size
        else:
            labels = None

        self.iter += 1
        return images, labels
//  [python]
def preprocess_image(self, batched_inputs, training):
        """
        Normalize, pad and batch the input images.
        """
        images = [x["image"].to(self.device) for x in batched_inputs]
        bs = len(images)
        images = [self.normalizer(x) for x in images]

        images = ImageList.from_tensors(
            images, size_divisibility=0, pad_ref_long=True)

        # sync image size for all gpus
        comm.synchronize()
        if training and self.iter % self.change_iter == 0:
            if self.iter < self.max_iter - 20000:
                meg = torch.LongTensor(1).to(self.device)
                comm.synchronize()
                if comm.is_main_process():
                    size = np.random.choice(self.multi_size)
                    meg.fill_(size)

                if comm.get_world_size() > 1:
                    comm.synchronize()
                    dist.broadcast(meg, 0)
                self.size = meg.item()

                comm.synchronize()
            else:
                self.size = 608

        if training:

            # resize image inputs
            modes = ['bilinear', 'nearest', 'bicubic', 'area']
            mode = modes[random.randrange(4)]
            if mode == 'bilinear' or mode == 'bicubic':
                images.tensor = F.interpolate(
                    images.tensor, size=[self.size, self.size], mode=mode, align_corners=False)
            else:
                images.tensor = F.interpolate(images.tensor, size=[self.size, self.size], mode=mode)

            if "instances" in batched_inputs[0]:
                gt_instances = [
                    x["instances"].to(self.device) for x in batched_inputs
                ]
            elif "targets" in batched_inputs[0]:
                log_first_n(
                    logging.WARN,
                    "'targets' in the model inputs is now renamed to 'instances'!",
                    n=10)
                gt_instances = [
                    x["targets"].to(self.device) for x in batched_inputs
                ]
            else:
                gt_instances = None

            targets = [
                torch.cat(
                    [instance.gt_classes.float().unsqueeze(-1), instance.gt_boxes.tensor], dim=-1
                )
                for instance in gt_instances
            ]
            labels = torch.zeros((bs, 100, 5))
            for i, target in enumerate(targets):
                labels[i][:target.shape[0]] = target
            labels[:, :, 1:] = labels[:, :, 1:] / 512. * self.size
        else:
            labels = None

        self.iter += 1
        return images, labels
//  [python]
def permute_diar(self, waves, spk_prediction):
        # Permute the diarization result using the correlation
        # between wav and spk_prediction
        # FIXME(YushiUeda): batch_size > 1 is not considered
        num_spk = len(waves)
        permute_list = [np.array(p) for p in permutations(range(num_spk))]
        corr_list = []
        interp_prediction = t.nn.functional.interpolate(
            torch.sigmoid(spk_prediction).transpose(1, 2),
            size=waves[0].size(1),
            mode="linear",
        ).transpose(1, 2)
        for p in permute_list:
            diar_perm = interp_prediction[:, :, p]
            corr_perm = [0]
            for q in range(num_spk):
                corr_perm += np.corrcoef(
                    torch.squeeze(abs(waves[q])).cpu().numpy(),
                    torch.squeeze(diar_perm[:, :, q]).cpu().numpy(),
                )[0, 1]
            corr_list.append(corr_perm)
        max_corr, max_idx = torch.max(torch.from_numpy(np.array(corr_list)), dim=0)
        return (
            spk_prediction[:, :, permute_list[max_idx]],
            interp_prediction[:, :, permute_list[max_idx]],
            permute_list[max_idx],
        )
//  [python]
def permute_diar(self, waves, spk_prediction):
        # Permute the diarization result using the correlation
        # between wav and spk_prediction
        # FIXME(YushiUeda): batch_size > 1 is not considered
        num_spk = len(waves)
        permute_list = [np.array(p) for p in permutations(range(num_spk))]
        corr_list = []
        interp_prediction = F.interpolate(
            torch.sigmoid(spk_prediction).transpose(1, 2),
            size=waves[0].size(1),
            mode="linear",
        ).transpose(1, 2)
        for p in permute_list:
            diar_perm = interp_prediction[:, :, p]
            corr_perm = [0]
            for q in range(num_spk):
                corr_perm += np.corrcoef(
                    torch.squeeze(abs(waves[q])).cpu().numpy(),
                    torch.squeeze(diar_perm[:, :, q]).cpu().numpy(),
                )[0, 1]
            corr_list.append(corr_perm)
        max_corr, max_idx = torch.max(torch.from_numpy(np.array(corr_list)), dim=0)
        return (
            spk_prediction[:, :, permute_list[max_idx]],
            interp_prediction[:, :, permute_list[max_idx]],
            permute_list[max_idx],
        )
//  [python]
def restore_size(self, pred_alpha: torch.Tensor,
                     data_sample: DataSample) -> torch.Tensor:
        """Restore the predicted alpha to the original shape.

        The shape of the predicted alpha may not be the same as the shape of
        original input image. This function restores the shape of the predicted
        alpha.

        Args:
            pred_alpha (torch.Tensor): A single predicted alpha of
                shape (1, H, W).
            data_sample (DataSample): Data sample containing
                original shape as meta data.

        Returns:
            torch.Tensor: The reshaped predicted alpha.
        """
        resize_method = self.test_cfg['resize_method']
        resize_mode = self.test_cfg['resize_mode']

        ori_h, ori_w = data_sample.ori_merged_shape[:2]
        if resize_method == 'pad':
            pred_alpha = pred_alpha[:, :ori_h, :ori_w]
        elif resize_method == 'interp':
            pred_alpha = t.nn.functional.interpolate(
                pred_alpha.unsqueeze(0), size=(ori_h, ori_w), mode=resize_mode)
            pred_alpha = pred_alpha[0]  # 1,H,W

        return pred_alpha
//  [python]
def restore_size(self, pred_alpha: torch.Tensor,
                     data_sample: DataSample) -> torch.Tensor:
        """Restore the predicted alpha to the original shape.

        The shape of the predicted alpha may not be the same as the shape of
        original input image. This function restores the shape of the predicted
        alpha.

        Args:
            pred_alpha (torch.Tensor): A single predicted alpha of
                shape (1, H, W).
            data_sample (DataSample): Data sample containing
                original shape as meta data.

        Returns:
            torch.Tensor: The reshaped predicted alpha.
        """
        resize_method = self.test_cfg['resize_method']
        resize_mode = self.test_cfg['resize_mode']

        ori_h, ori_w = data_sample.ori_merged_shape[:2]
        if resize_method == 'pad':
            pred_alpha = pred_alpha[:, :ori_h, :ori_w]
        elif resize_method == 'interp':
            pred_alpha = F.interpolate(
                pred_alpha.unsqueeze(0), size=(ori_h, ori_w), mode=resize_mode)
            pred_alpha = pred_alpha[0]  # 1,H,W

        return pred_alpha
//  [python]
def forward(self, x, concat_with):
        up_x = F.upsample_bilinear(
            x,
            size=[concat_with.size(2),
                  concat_with.size(3)])
        f = torch.cat([up_x, concat_with], dim=1)
        return self._net(f)
//  [python]
def forward(self, x, concat_with):
        up_x = F.interpolate(
            x,
            size=[concat_with.size(2),
                  concat_with.size(3)],
            mode='bilinear',
            align_corners=True)
        f = torch.cat([up_x, concat_with], dim=1)
        return self._net(f)
//  [python]
def forward(self, x):
    x = self.stem(x)
    x = self.layer1(x)

    x_list = []
    for i in range(self.config.S2_NUM_BRANCHES):
      if self.transition1[i] is not None:
        x_list.append(self.transition1[i](x))
      else:
        x_list.append(x)
    y_list = self.stage2(x_list)

    x_list = []
    for i in range(self.config.S3_NUM_BRANCHES):
      if self.transition2[i] is not None:
        x_list.append(self.transition2[i](y_list[-1]))
      else:
        x_list.append(y_list[i])
    y_list = self.stage3(x_list)

    x_list = []
    for i in range(self.config.S4_NUM_BRANCHES):
      if self.transition3[i] is not None:
        x_list.append(self.transition3[i](y_list[-1]))
      else:
        x_list.append(y_list[i])
    feats = self.stage4(x_list)

    # Head Part
    batch_size, _, height, width = feats[0].shape
    feat1 = t.nn.functional.interpolate(feats[1], size=(height, width), mode='bilinear', align_corners=False)
    feat2 = t.nn.functional.interpolate(feats[2], size=(height, width), mode='bilinear', align_corners=False)
    feat3 = t.nn.functional.interpolate(feats[3], size=(height, width), mode='bilinear', align_corners=False)
    cat_f = torch.cat([feats[0], feat1, feat2, feat3], 1)
    heatmaps = self.head(cat_f)

    # The location of the current batch
    batch_locs, batch_scos = [], []
    for ibatch in range(batch_size):
      batch_location, batch_score = find_tensor_peak_batch(heatmaps[ibatch], self.sigma, self.downsample)
      batch_locs.append( batch_location )
      batch_scos.append( batch_score )
    batch_locs, batch_scos = torch.stack(batch_locs), torch.stack(batch_scos)
    return feats, [heatmaps], batch_locs, batch_scos
//  [python]
def forward(self, x):
    x = self.stem(x)
    x = self.layer1(x)

    x_list = []
    for i in range(self.config.S2_NUM_BRANCHES):
      if self.transition1[i] is not None:
        x_list.append(self.transition1[i](x))
      else:
        x_list.append(x)
    y_list = self.stage2(x_list)

    x_list = []
    for i in range(self.config.S3_NUM_BRANCHES):
      if self.transition2[i] is not None:
        x_list.append(self.transition2[i](y_list[-1]))
      else:
        x_list.append(y_list[i])
    y_list = self.stage3(x_list)

    x_list = []
    for i in range(self.config.S4_NUM_BRANCHES):
      if self.transition3[i] is not None:
        x_list.append(self.transition3[i](y_list[-1]))
      else:
        x_list.append(y_list[i])
    feats = self.stage4(x_list)

    # Head Part
    batch_size, _, height, width = feats[0].shape
    feat1 = F.interpolate(feats[1], size=(height, width), mode='bilinear', align_corners=False)
    feat2 = F.interpolate(feats[2], size=(height, width), mode='bilinear', align_corners=False)
    feat3 = F.interpolate(feats[3], size=(height, width), mode='bilinear', align_corners=False)
    cat_f = torch.cat([feats[0], feat1, feat2, feat3], 1)
    heatmaps = self.head(cat_f)

    # The location of the current batch
    batch_locs, batch_scos = [], []
    for ibatch in range(batch_size):
      batch_location, batch_score = find_tensor_peak_batch(heatmaps[ibatch], self.sigma, self.downsample)
      batch_locs.append( batch_location )
      batch_scos.append( batch_score )
    batch_locs, batch_scos = torch.stack(batch_locs), torch.stack(batch_scos)
    return feats, [heatmaps], batch_locs, batch_scos
//  [python]
def interpolate_as(input, other, mode='nearest', align_corners=False):
    """interpolate_as"""
    if mode == 'nearest':
        align_corners = None
    return t.interpolate(input, size=other.size()[2:], mode=mode, align_corners=align_corners)
//  [python]
def interpolate_as(input, other, mode='nearest', align_corners=False):
    """interpolate_as"""
    if mode == 'nearest':
        align_corners = None
    return F.interpolate(input, size=other.size()[2:], mode=mode, align_corners=align_corners)
//  [python]
def forward(self, x):
        mini_size = x[-1].size()[-2:]
        out = [F.adaptive_avg_pool2d(s, mini_size) for s in x[:-1]] + [x[-1]]
        out = torch.cat(out, dim=1)
        out = self.conv1(out)
        out = self.conv2(out)
        out = torch.split(out, self.channels, dim=1)
        out = [
            s * F.upsample(a, size=s.size()[-2:], mode='nearest')
            for s, a in zip(x, out)
        ]
        return out
//  [python]
def forward(self, x):
        mini_size = x[-1].size()[-2:]
        out = [F.adaptive_avg_pool2d(s, mini_size) for s in x[:-1]] + [x[-1]]
        out = torch.cat(out, dim=1)
        out = self.conv1(out)
        out = self.conv2(out)
        out = torch.split(out, self.channels, dim=1)
        out = [
            s * F.interpolate(a, size=s.size()[-2:], mode='nearest')
            for s, a in zip(x, out)
        ]
        return out
//  [python]
def forward(self, x):
        x = x[::-1]

        y = []
        last_x = None
        for i, s in enumerate(x):
            if last_x is not None:
                last_x = t.interpolate(
                    last_x,
                    size=s.size()[-2:],
                    mode='bilinear',
                    align_corners=True)
                s = s + last_x
            s = self.projects[i](s)
            y.append(s)
            last_x = s

        return y[::-1]
//  [python]
def forward(self, x):
        x = x[::-1]

        y = []
        last_x = None
        for i, s in enumerate(x):
            if last_x is not None:
                last_x = F.interpolate(
                    last_x,
                    size=s.size()[-2:],
                    mode='bilinear',
                    align_corners=True)
                s = s + last_x
            s = self.projects[i](s)
            y.append(s)
            last_x = s

        return y[::-1]
//  [python]
def forward(self, x):
        B, C, N = x.shape
        x = x + F.upsample(
            self.pos_embed, size=(N, ), mode='linear')
        return x
//  [python]
def forward(self, x):
        B, C, N = x.shape
        x = x + F.interpolate(
            self.pos_embed, size=(N, ), mode='linear', align_corners=False)
        return x
//  [python]
def forward(self, x_l, x_g):
        """
        x_g: global features
        x_l: local features
        """
        B, C, H, W = x_l.shape

        local_feat = self.local_embedding(x_l)
        global_act = self.global_act(x_g)
        sig_act = t.interpolate(
            global_act, size=(H, W), mode='bilinear', align_corners=False)

        out = local_feat * sig_act

        return out
//  [python]
def forward(self, x_l, x_g):
        """
        x_g: global features
        x_l: local features
        """
        B, C, H, W = x_l.shape

        local_feat = self.local_embedding(x_l)
        global_act = self.global_act(x_g)
        sig_act = F.interpolate(
            global_act, size=(H, W), mode='bilinear', align_corners=False)

        out = local_feat * sig_act

        return out
//  [python]
def forward(self, x):
        x = (x + 1) / 2
        x = F.interpolate(x, self.size, mode='bicubic', align_corners=False, antialias=True)
        if x.shape[1] == 1:
            x = torch.cat([x] * 3, dim=1)
        x = self.normalize(x)
        x = self.model.encode_image(x).float()
        x = F.normalize(x) * x.shape[-1] ** 0.5
        return x
//  [python]
def forward(self, inputs):
        # x_x8, x_x16, x_x32
        low_feat1 = self.embedding_list[0](inputs[0])

        low_feat2 = t.interpolate(
            inputs[1], size=low_feat1.shape[-2:], mode='bilinear')
        low_feat2_act = self.act_embedding_list[0](low_feat2)

        high_feat_act = t.interpolate(
            self.act_embedding_list[1](inputs[2]),
            size=low_feat2.shape[2:],
            mode='bilinear',
        )
        high_feat = t.interpolate(
            self.embedding_list[1](inputs[2]),
            size=low_feat2.shape[2:],
            mode='bilinear')

        res = low_feat2_act * high_feat_act * low_feat1 + high_feat

        return res
//  [python]
def forward(self, inputs):
        # x_x8, x_x16, x_x32
        low_feat1 = self.embedding_list[0](inputs[0])

        low_feat2 = F.interpolate(
            inputs[1], size=low_feat1.shape[-2:], mode='bilinear')
        low_feat2_act = self.act_embedding_list[0](low_feat2)

        high_feat_act = F.interpolate(
            self.act_embedding_list[1](inputs[2]),
            size=low_feat2.shape[2:],
            mode='bilinear',
        )
        high_feat = F.interpolate(
            self.embedding_list[1](inputs[2]),
            size=low_feat2.shape[2:],
            mode='bilinear')

        res = low_feat2_act * high_feat_act * low_feat1 + high_feat

        return res
//  [python]
def forward(self, inputs):  # x_x8, x_x16, x_x32, x_x64
        low_feat1 = t.interpolate(inputs[0], scale_factor=0.5, mode='bilinear')
        low_feat1_act = self.act_embedding_list[0](low_feat1)
        low_feat1 = self.embedding_list[0](low_feat1)

        low_feat2 = t.interpolate(
            inputs[1], size=low_feat1.shape[-2:], mode='bilinear')
        low_feat2_act = self.act_embedding_list[1](low_feat2)  # x16
        low_feat2 = self.embedding_list[1](low_feat2)

        high_feat_act = t.interpolate(
            self.act_embedding_list[2](inputs[2]),
            size=low_feat2.shape[2:],
            mode='bilinear',
        )
        high_feat = t.interpolate(
            self.embedding_list[2](inputs[2]),
            size=low_feat2.shape[2:],
            mode='bilinear')

        res = (
            low_feat1_act * low_feat2_act * high_feat_act *
            (low_feat1 + low_feat2) + high_feat)

        return res