jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x):
        H, W = x.size()[2:]

        x = self.init_block(x)

        x1 = self.block1_1(x)
        x1 = self.block1_2(x1)
        x1 = self.block1_3(x1)

        x2 = self.block2_1(x1)
        x2 = self.block2_2(x2)
        x2 = self.block2_3(x2)
        x2 = self.block2_4(x2)

        x3 = self.block3_1(x2)
        x3 = self.block3_2(x3)
        x3 = self.block3_3(x3)
        x3 = self.block3_4(x3)

        x4 = self.block4_1(x3)
        x4 = self.block4_2(x4)
        x4 = self.block4_3(x4)
        x4 = self.block4_4(x4)

        x_fuses = []
        if self.sa and self.dil is not None:
            for i, xi in enumerate([x1, x2, x3, x4]):
                x_fuses.append(self.attentions[i](self.dilations[i](xi)))
        elif self.sa:
            for i, xi in enumerate([x1, x2, x3, x4]):
                x_fuses.append(self.attentions[i](xi))
        elif self.dil is not None:
            for i, xi in enumerate([x1, x2, x3, x4]):
                x_fuses.append(self.dilations[i](xi))
        else:
            x_fuses = [x1, x2, x3, x4]

        e1 = self.conv_reduces[0](x_fuses[0])
        e1 = F.interpolate(e1, (H, W), mode="bilinear", align_corners=False)

        e2 = self.conv_reduces[1](x_fuses[1])
        e2 = F.interpolate(e2, (H, W), mode="bilinear", align_corners=False)

        e3 = self.conv_reduces[2](x_fuses[2])
        e3 = F.interpolate(e3, (H, W), mode="bilinear", align_corners=False)

        e4 = self.conv_reduces[3](x_fuses[3])
        e4 = F.interpolate(e4, (H, W), mode="bilinear", align_corners=False)

        outputs = [e1, e2, e3, e4]

        output = self.classifier(torch.cat(outputs, dim=1))
        # if not self.training:
        #    return torch.sigmoid(output)

        outputs.append(output)
        outputs = [torch.sigmoid(r) for r in outputs]
        return outputs
//  [python]
def forward(self, x):
        for i, f in enumerate(x):
            if i == 0:
                r = self.refine[i](f)
            else:
                r_p = self.refine[i](f)
                r_p = t.interpolate(r_p, r.size()[2:], mode="bilinear", align_corners=False)
                r = r + r_p
        return r
//  [python]
def forward(self, x):
        for i, f in enumerate(x):
            if i == 0:
                r = self.refine[i](f)
            else:
                r_p = self.refine[i](f)
                r_p = F.interpolate(r_p, r.size()[2:], mode="bilinear", align_corners=False)
                r = r + r_p
        return r
//  [python]
def post_process(
    input_array,
    attn_masks,
    pad=None,
    mask_threshold=0.3,
    return_largest=False,
    min_area_ratio=0.15,
    return_instances=False,
):
  """post process the input tensor with the attention masks.

  Args:
      input_array: A np.ndarray input array to be post processed with shape
        [width, height, 3, batch_size]
      attn_masks: A torch.Tensor for the attention masks with shape [1,
        num_texts, width, height]
      pad: A list of padding: [pad_left, pad_top, width, height], where
        pad_left, pad_top and width, height are int values.
      mask_threshold: The threshold to binarize the mask.
      return_largest: If true, return the largest connected component.
      min_area_ratio: Keep the mask if its area is larger than this threshold.
      return_instances: Whether to return instances or not.

  Returns:
      attn_masks: A list of tensors with shape [num_instances, height, width]
          x num_texts, where len(attn_masks) = num_texts.
      NOTE: the number_instances for each text (class) may vary.
      The output is a binary tensor.
  """
  if len(attn_masks.shape) == 3:
    attn_masks = attn_masks[None]
  img_width, img_height = input_array.shape[:2]
  attn_masks = t.nn.functional.interpolate(
      attn_masks, size=(img_height, img_width), mode='bicubic'
  ).squeeze(0)
  device = attn_masks.device
  output_masks = filter_masks(
      attn_masks,
      pad=pad,
      mask_threshold=mask_threshold,
      min_area_ratio=min_area_ratio,
      return_largest=return_largest,
      device=device,
      return_instances=return_instances,
  )
  if pad is not None:
    left, top, width, height = pad
    input_array = input_array[top : top + height, left : left + width]
  return input_array, output_masks
//  [python]
def post_process(
    input_array,
    attn_masks,
    pad=None,
    mask_threshold=0.3,
    return_largest=False,
    min_area_ratio=0.15,
    return_instances=False,
):
  """post process the input tensor with the attention masks.

  Args:
      input_array: A np.ndarray input array to be post processed with shape
        [width, height, 3, batch_size]
      attn_masks: A torch.Tensor for the attention masks with shape [1,
        num_texts, width, height]
      pad: A list of padding: [pad_left, pad_top, width, height], where
        pad_left, pad_top and width, height are int values.
      mask_threshold: The threshold to binarize the mask.
      return_largest: If true, return the largest connected component.
      min_area_ratio: Keep the mask if its area is larger than this threshold.
      return_instances: Whether to return instances or not.

  Returns:
      attn_masks: A list of tensors with shape [num_instances, height, width]
          x num_texts, where len(attn_masks) = num_texts.
      NOTE: the number_instances for each text (class) may vary.
      The output is a binary tensor.
  """
  if len(attn_masks.shape) == 3:
    attn_masks = attn_masks[None]
  img_width, img_height = input_array.shape[:2]
  attn_masks = F.interpolate(
      attn_masks, size=(img_height, img_width), mode='bicubic'
  ).squeeze(0)
  device = attn_masks.device
  output_masks = filter_masks(
      attn_masks,
      pad=pad,
      mask_threshold=mask_threshold,
      min_area_ratio=min_area_ratio,
      return_largest=return_largest,
      device=device,
      return_instances=return_instances,
  )
  if pad is not None:
    left, top, width, height = pad
    input_array = input_array[top : top + height, left : left + width]
  return input_array, output_masks
//  [python]
def load_pretrained(self, state_dict):
        
        new_state_dict = {}
        for state_key, state_value in state_dict.items():
            keys = state_key.split('.')
            m = self
            for key in keys:
                if key.isdigit():
                    m = m[int(key)]
                else:
                    m = getattr(m, key)
            if m.shape == state_value.shape:
                new_state_dict[state_key] = state_value
            else:
                # Ignore different shapes
                if 'relative_position_index' in keys:
                    new_state_dict[state_key] = m.data
                if 'q_grid' in keys:
                    new_state_dict[state_key] = m.data
                if 'reference' in keys:
                    new_state_dict[state_key] = m.data
                # Bicubic Interpolation
                if 'relative_position_bias_table' in keys:
                    n, c = state_value.size()
                    l = int(math.sqrt(n))
                    assert n == l ** 2
                    L = int(math.sqrt(m.shape[0]))
                    pre_interp = state_value.reshape(1, l, l, c).permute(0, 3, 1, 2)
                    post_interp = t.interpolate(pre_interp, (L, L), mode='bicubic')
                    new_state_dict[state_key] = post_interp.reshape(c, L ** 2).permute(1, 0)
                if 'rpe_table' in keys:
                    c, h, w = state_value.size()
                    C, H, W = m.data.size()
                    pre_interp = state_value.unsqueeze(0)
                    post_interp = t.interpolate(pre_interp, (H, W), mode='bicubic')
                    new_state_dict[state_key] = post_interp.squeeze(0)
        
        self.load_state_dict(new_state_dict, strict=False)
//  [python]
def load_pretrained(self, state_dict):
        
        new_state_dict = {}
        for state_key, state_value in state_dict.items():
            keys = state_key.split('.')
            m = self
            for key in keys:
                if key.isdigit():
                    m = m[int(key)]
                else:
                    m = getattr(m, key)
            if m.shape == state_value.shape:
                new_state_dict[state_key] = state_value
            else:
                # Ignore different shapes
                if 'relative_position_index' in keys:
                    new_state_dict[state_key] = m.data
                if 'q_grid' in keys:
                    new_state_dict[state_key] = m.data
                if 'reference' in keys:
                    new_state_dict[state_key] = m.data
                # Bicubic Interpolation
                if 'relative_position_bias_table' in keys:
                    n, c = state_value.size()
                    l = int(math.sqrt(n))
                    assert n == l ** 2
                    L = int(math.sqrt(m.shape[0]))
                    pre_interp = state_value.reshape(1, l, l, c).permute(0, 3, 1, 2)
                    post_interp = F.interpolate(pre_interp, (L, L), mode='bicubic')
                    new_state_dict[state_key] = post_interp.reshape(c, L ** 2).permute(1, 0)
                if 'rpe_table' in keys:
                    c, h, w = state_value.size()
                    C, H, W = m.data.size()
                    pre_interp = state_value.unsqueeze(0)
                    post_interp = F.interpolate(pre_interp, (H, W), mode='bicubic')
                    new_state_dict[state_key] = post_interp.squeeze(0)
        
        self.load_state_dict(new_state_dict, strict=False)
//  [python]
def _compute_distance(self, images):
        """Feed data to the metric.

        Args:
            images (Tensor): Input tensor.
        """
        # use minibatch's device type to initialize a lpips calculator
        if not hasattr(self, 'percept'):
            self.percept = lpips.LPIPS(net='vgg').to(images.device)
        # crop and resize images
        if self.crop:
            c = images.shape[2] // 8
            minibatch = images[:, :, c * 3:c * 7, c * 2:c * 6]

        factor = minibatch.shape[2] // 256
        if factor > 1:
            minibatch = t.nn.functional.interpolate(
                minibatch,
                size=(256, 256),
                mode='bilinear',
                align_corners=False)
        # calculator and store lpips score
        distance = self.percept(minibatch[::2], minibatch[1::2]).view(
            minibatch.shape[0] // 2) / (
                self.epsilon**2)
        return distance.to('cpu')
//  [python]
def _compute_distance(self, images):
        """Feed data to the metric.

        Args:
            images (Tensor): Input tensor.
        """
        # use minibatch's device type to initialize a lpips calculator
        if not hasattr(self, 'percept'):
            self.percept = lpips.LPIPS(net='vgg').to(images.device)
        # crop and resize images
        if self.crop:
            c = images.shape[2] // 8
            minibatch = images[:, :, c * 3:c * 7, c * 2:c * 6]

        factor = minibatch.shape[2] // 256
        if factor > 1:
            minibatch = F.interpolate(
                minibatch,
                size=(256, 256),
                mode='bilinear',
                align_corners=False)
        # calculator and store lpips score
        distance = self.percept(minibatch[::2], minibatch[1::2]).view(
            minibatch.shape[0] // 2) / (
                self.epsilon**2)
        return distance.to('cpu')
//  [python]
def forward(self, inputs, logdets):
        x0, x1 = inputs
        logdet0, logdet1 = logdets
        if self.direct == 'up':
            y10 = t.interpolate(x1, size=x0.shape[2:], mode='nearest') # interpolation first in up-sampling
            (scale0, bias0), dlogdet0 = self.affineParams(y10)
            z0, z1 = scale0*x0+bias0, x1
            dlogdet1 = 0
        else:
            (scale10, bias10), dlogdet10 = self.affineParams(x0)
            scale1, bias1, dlogdet1 = t.interpolate(scale10, size=x1.shape[2:], mode='nearest'),\
                                         t.interpolate(bias10, size=x1.shape[2:], mode='nearest'),\
                                             t.interpolate(dlogdet10, size=x1.shape[2:], mode='nearest') # interpolation after in down-sampling
            z0, z1 = x0, scale1*x1+bias1
            dlogdet0 = 0
        outputs = (z0, z1)
        out_logdets = (logdet0+dlogdet0, logdet1+dlogdet1)
        return outputs, out_logdets
//  [python]
def forward(self, inputs, logdets):
        x0, x1 = inputs
        logdet0, logdet1 = logdets
        if self.direct == 'up':
            y10 = F.interpolate(x1, size=x0.shape[2:], mode='nearest') # interpolation first in up-sampling
            (scale0, bias0), dlogdet0 = self.affineParams(y10)
            z0, z1 = scale0*x0+bias0, x1
            dlogdet1 = 0
        else:
            (scale10, bias10), dlogdet10 = self.affineParams(x0)
            scale1, bias1, dlogdet1 = F.interpolate(scale10, size=x1.shape[2:], mode='nearest'),\
                                         F.interpolate(bias10, size=x1.shape[2:], mode='nearest'),\
                                             F.interpolate(dlogdet10, size=x1.shape[2:], mode='nearest') # interpolation after in down-sampling
            z0, z1 = x0, scale1*x1+bias1
            dlogdet0 = 0
        outputs = (z0, z1)
        out_logdets = (logdet0+dlogdet0, logdet1+dlogdet1)
        return outputs, out_logdets
//  [python]
def forward(self, x):  # l,m,s
        """
        # 128, 256, 512
        512, 256, 128
        from small -> large
        """
        x_level_0 = x[2]  # l
        x_level_1 = x[1]  # m
        x_level_2 = x[0]  # s
        if self.level == 0:
            level_0_resized = x_level_0
            level_1_resized = self.stride_level_1(x_level_1)
            level_2_downsampled_inter = F.max_pool2d(x_level_2, 3, stride=2, padding=1)
            level_2_resized = self.stride_level_2(level_2_downsampled_inter)
        elif self.level == 1:
            level_0_compressed = self.compress_level_0(x_level_0)
            level_0_resized = F.interpolate(level_0_compressed, scale_factor=2, mode='nearest')
            level_1_resized = x_level_1
            level_2_resized = self.stride_level_2(x_level_2)
        elif self.level == 2:
            level_0_compressed = self.compress_level_0(x_level_0)
            level_0_resized = F.interpolate(level_0_compressed, scale_factor=4, mode='nearest')
            x_level_1_compressed = self.compress_level_1(x_level_1)
            level_1_resized = F.interpolate(x_level_1_compressed, scale_factor=2, mode='nearest')
            level_2_resized = x_level_2

        level_0_weight_v = self.weight_level_0(level_0_resized)
        level_1_weight_v = self.weight_level_1(level_1_resized)
        level_2_weight_v = self.weight_level_2(level_2_resized)

        levels_weight_v = torch.cat((level_0_weight_v, level_1_weight_v, level_2_weight_v), 1)
        levels_weight = self.weight_levels(levels_weight_v)
        levels_weight = F.softmax(levels_weight, dim=1)

        fused_out_reduced = level_0_resized * levels_weight[:, 0:1, :, :] + \
                            level_1_resized * levels_weight[:, 1:2, :, :] + \
                            level_2_resized * levels_weight[:, 2:, :, :]

        out = self.expand(fused_out_reduced)

        if self.vis:
            return out, levels_weight, fused_out_reduced.sum(dim=1)
        else:
            return out
//  [python]
def _pyramid_up(self, input, size, pad_mode='reflect'):
        if not len(input.shape) == 4:
            raise ValueError(f'Invalid img shape, we expect BCHW, got: {input.shape}')
        # upsample
        img_up = t.interpolate(input, size=size, mode='nearest', )
        # blur
        img_pad = F.pad(img_up, (2,2,2,2), mode=pad_mode)
        img_blur = kornia_filter2d(img_pad, kernel=self.kernel)
        return img_blur
//  [python]
def _pyramid_up(self, input, size, pad_mode='reflect'):
        if not len(input.shape) == 4:
            raise ValueError(f'Invalid img shape, we expect BCHW, got: {input.shape}')
        # upsample
        img_up = F.interpolate(input, size=size, mode='nearest', )
        # blur
        img_pad = F.pad(img_up, (2,2,2,2), mode=pad_mode)
        img_blur = kornia_filter2d(img_pad, kernel=self.kernel)
        return img_blur
//  [python]
def _pyramid_down(self, input, pad_mode='reflect'):
        if not len(input.shape) == 4:
            raise ValueError(f'Invalid img shape, we expect BCHW, got: {input.shape}')
        # blur 
        img_pad = F.pad(input, (2,2,2,2), mode=pad_mode)
        img_blur = kornia_filter2d(img_pad, kernel=self.kernel)
        # downsample
        out = t.nn.functional.interpolate(img_blur, scale_factor=0.5, mode='nearest')
        return out
//  [python]
def _pyramid_down(self, input, pad_mode='reflect'):
        if not len(input.shape) == 4:
            raise ValueError(f'Invalid img shape, we expect BCHW, got: {input.shape}')
        # blur 
        img_pad = F.pad(input, (2,2,2,2), mode=pad_mode)
        img_blur = kornia_filter2d(img_pad, kernel=self.kernel)
        # downsample
        out = F.interpolate(img_blur, scale_factor=0.5, mode='nearest')
        return out
//  [python]
def inverse(self, outputs, logdets):
        z0, z1 = outputs
        logdet0, logdet1 = logdets
        if self.direct == 'up':
            z10 = t.interpolate(z1, size=z0.shape[2:], mode='nearest') # interpolation first in up-sampling
            (scale0, bias0), dlogdet0 = self.affineParams(z10)
            x0, x1 = (z0-bias0)/scale0, z1
            dlogdet1 = 0
        else:
            (scale01, bias01), dlogdet01 = self.affineParams(z0)
            scale1, bias1, dlogdet1 = t.interpolate(scale01, size=z1.shape[2:], mode='nearest'),\
                                         t.interpolate(bias01, size=z1.shape[2:], mode='nearest'),\
                                             t.interpolate(dlogdet01, size=z1.shape[2:], mode='nearest') # interpolation after in down-sampling
            x0, x1 = z0, (z1-bias1)/scale1
            dlogdet0 = 0
        inputs = (x0, x1)
        in_logdets = (logdet0-dlogdet0, logdet1-dlogdet1)
        return inputs, in_logdets
//  [python]
def inverse(self, outputs, logdets):
        z0, z1 = outputs
        logdet0, logdet1 = logdets
        if self.direct == 'up':
            z10 = F.interpolate(z1, size=z0.shape[2:], mode='nearest') # interpolation first in up-sampling
            (scale0, bias0), dlogdet0 = self.affineParams(z10)
            x0, x1 = (z0-bias0)/scale0, z1
            dlogdet1 = 0
        else:
            (scale01, bias01), dlogdet01 = self.affineParams(z0)
            scale1, bias1, dlogdet1 = F.interpolate(scale01, size=z1.shape[2:], mode='nearest'),\
                                         F.interpolate(bias01, size=z1.shape[2:], mode='nearest'),\
                                             F.interpolate(dlogdet01, size=z1.shape[2:], mode='nearest') # interpolation after in down-sampling
            x0, x1 = z0, (z1-bias1)/scale1
            dlogdet0 = 0
        inputs = (x0, x1)
        in_logdets = (logdet0-dlogdet0, logdet1-dlogdet1)
        return inputs, in_logdets
//  [python]
def forward(self, pred, out):
        """
        :param pred:
        :param out:
        :return:
        """
        ASFF = False  # https://arxiv.org/abs/1911.09516
        if ASFF:
            i, n = self.inde
            x, self.nl  # index in layers, number of layers
            pred = out[self.layers[i]]
            bs, _, ny, nx = pred.shape  # bs, 255, 13, 13
            if (self.nx, self.ny) != (nx, ny):
                self.create_grids((nx, ny), pred.device)

            # outputs and weights
            # w = F.softmax(p[:, -n:], 1)  # normalized weights
            w = torch.sigmoid(pred[:, -n:]) * (2 / n)  # sigmoid weights (faster)
            # w = w / w.sum(1).unsqueeze(1)  # normalize across layer dimension

            # weighted ASFF sum
            pred = out[self.layers[i]][:, :-n] * w[:, i:i + 1]
            for j in range(n):
                if j != i:
                    pred += w[:, j:j + 1] * \
                            t.interpolate(out[self.layers[j]][:, :-n], size=[ny, nx], mode='bilinear',
                                          align_corners=False)

        else:
            bs, _, ny, nx = pred.shape  # bs, 255, 13, 13
            if (self.nx, self.ny) != (nx, ny):
                self.create_grids(ng=(nx, ny), device=pred.device)

        # p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85)  # (bs, na, ny, nx, no(classes + xywh))
        pred = pred.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous()  # prediction

        if self.training:
            return pred

        else:  # inference
            io = pred.clone()  # inference output

            # ---------- process pred to io
            io[..., :2] = torch.sigmoid(io[..., :2]) + self.grid  # xy
            io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh  # wh YOLO method
            io[..., :4] *= self.stride  # map from YOLO layer's scale to net input's scale
            torch.sigmoid_(io[..., 4:])  # sigmoid for confidence score and cls pred

            # gathered pred output: io: view [1, 3, 13, 13, 85] as [1, 507, 85]
            io = io.view(bs, -1, self.no)

            # return io, pred
            return io, pred
//  [python]
def forward(self, pred, out):
        """
        :param pred:
        :param out:
        :return:
        """
        ASFF = False  # https://arxiv.org/abs/1911.09516
        if ASFF:
            i, n = self.inde
            x, self.nl  # index in layers, number of layers
            pred = out[self.layers[i]]
            bs, _, ny, nx = pred.shape  # bs, 255, 13, 13
            if (self.nx, self.ny) != (nx, ny):
                self.create_grids((nx, ny), pred.device)

            # outputs and weights
            # w = F.softmax(p[:, -n:], 1)  # normalized weights
            w = torch.sigmoid(pred[:, -n:]) * (2 / n)  # sigmoid weights (faster)
            # w = w / w.sum(1).unsqueeze(1)  # normalize across layer dimension

            # weighted ASFF sum
            pred = out[self.layers[i]][:, :-n] * w[:, i:i + 1]
            for j in range(n):
                if j != i:
                    pred += w[:, j:j + 1] * \
                            F.interpolate(out[self.layers[j]][:, :-n], size=[ny, nx], mode='bilinear',
                                          align_corners=False)

        else:
            bs, _, ny, nx = pred.shape  # bs, 255, 13, 13
            if (self.nx, self.ny) != (nx, ny):
                self.create_grids(ng=(nx, ny), device=pred.device)

        # p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85)  # (bs, na, ny, nx, no(classes + xywh))
        pred = pred.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous()  # prediction

        if self.training:
            return pred

        else:  # inference
            io = pred.clone()  # inference output

            # ---------- process pred to io
            io[..., :2] = torch.sigmoid(io[..., :2]) + self.grid  # xy
            io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh  # wh YOLO method
            io[..., :4] *= self.stride  # map from YOLO layer's scale to net input's scale
            torch.sigmoid_(io[..., 4:])  # sigmoid for confidence score and cls pred

            # gathered pred output: io: view [1, 3, 13, 13, 85] as [1, 507, 85]
            io = io.view(bs, -1, self.no)

            # return io, pred
            return io, pred
//  [python]
def forward(self, x):
        num_upsample = self.num_upsample
        for i in range(self.num_layers):
            x = self.conv[i](x)
            if num_upsample > 0:
                num_upsample -= 1
                x = F.upsample(
                    x, scale_factor=2, mode='bilinear')
        return x
//  [python]
def forward(self, x):
        num_upsample = self.num_upsample
        for i in range(self.num_layers):
            x = self.conv[i](x)
            if num_upsample > 0:
                num_upsample -= 1
                x = F.interpolate(
                    x, scale_factor=2, mode='bilinear', align_corners=False)
        return x
//  [python]
def test_schema_human_readable(self):
        """
        Make sure that the schema is human readable, ie the mode parameter should read "nearest" instead of being displayed in octal
        aten::__interpolate(Tensor input, int? size=None, float[]? scale_factor=None,
        str mode='\156\145\141\162\145\163\164', bool? align_corners=None) -> (Tensor):
        Expected a value of type 'Optional[int]' for argument 'size' but instead found type 'Tensor'.
        """
        with self.assertRaisesRegexWithHighlight(RuntimeError, "nearest", ""):

            @torch.jit.script
            def FooTest(x):
                return t.interpolate(x, "bad")
//  [python]
def test_schema_human_readable(self):
        """
        Make sure that the schema is human readable, ie the mode parameter should read "nearest" instead of being displayed in octal
        aten::__interpolate(Tensor input, int? size=None, float[]? scale_factor=None,
        str mode='\156\145\141\162\145\163\164', bool? align_corners=None) -> (Tensor):
        Expected a value of type 'Optional[int]' for argument 'size' but instead found type 'Tensor'.
        """
        with self.assertRaisesRegexWithHighlight(RuntimeError, "nearest", ""):

            @torch.jit.script
            def FooTest(x):
                return torch.nn.functional.interpolate(x, "bad")
//  [python]
def forward(self, x, segmap):

        # Part 1. generate parameter-free normalized activations
        normalized = self.param_free_norm(x)

        # Part 2. produce scaling and bias conditioned on semantic map
        segmap = t.interpolate(segmap, size=x.size()[2:], mode='nearest')
        actv = self.mlp_shared(segmap)
        gamma = self.mlp_gamma(actv)
        beta = self.mlp_beta(actv)

        # apply scale and bias
        out = normalized * (1 + gamma) + beta

        return out
//  [python]
def forward(self, x, segmap):

        # Part 1. generate parameter-free normalized activations
        normalized = self.param_free_norm(x)

        # Part 2. produce scaling and bias conditioned on semantic map
        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
        actv = self.mlp_shared(segmap)
        gamma = self.mlp_gamma(actv)
        beta = self.mlp_beta(actv)

        # apply scale and bias
        out = normalized * (1 + gamma) + beta

        return out
//  [python]
def upsampling_z(self, z, slice_ids=None, y_lengths=None, y_mask=None):
        spec_segment_size = self.spec_segment_size
        if self.args.encoder_sample_rate:
            # recompute the slices and spec_segment_size if needed
            slice_ids = slice_ids * int(self.interpolate_factor) if slice_ids is not None else slice_ids
            spec_segment_size = spec_segment_size * int(self.interpolate_factor)
            # interpolate z if needed
            if self.args.interpolate_z:
                z = t.nn.functional.interpolate(z, scale_factor=[self.interpolate_factor], mode="linear").squeeze(0)
                # recompute the mask if needed
                if y_lengths is not None and y_mask is not None:
                    y_mask = (
                        sequence_mask(y_lengths * self.interpolate_factor, None).to(y_mask.dtype).unsqueeze(1)
                    )  # [B, 1, T_dec_resampled]

        return z, spec_segment_size, slice_ids, y_mask
//  [python]
def upsampling_z(self, z, slice_ids=None, y_lengths=None, y_mask=None):
        spec_segment_size = self.spec_segment_size
        if self.args.encoder_sample_rate:
            # recompute the slices and spec_segment_size if needed
            slice_ids = slice_ids * int(self.interpolate_factor) if slice_ids is not None else slice_ids
            spec_segment_size = spec_segment_size * int(self.interpolate_factor)
            # interpolate z if needed
            if self.args.interpolate_z:
                z = torch.nn.functional.interpolate(z, scale_factor=[self.interpolate_factor], mode="linear").squeeze(0)
                # recompute the mask if needed
                if y_lengths is not None and y_mask is not None:
                    y_mask = (
                        sequence_mask(y_lengths * self.interpolate_factor, None).to(y_mask.dtype).unsqueeze(1)
                    )  # [B, 1, T_dec_resampled]

        return z, spec_segment_size, slice_ids, y_mask
//  [python]
def split_feats(self, feats):
        return (t.interpolate(feats[0], scale_factor=0.5, mode='bilinear'),
                feats[1],
                feats[2],
                feats[3],
                t.interpolate(feats[4], size=feats[3].shape[-2:], mode='bilinear'))
//  [python]
def split_feats(self, feats):
        return (F.interpolate(feats[0], scale_factor=0.5, mode='bilinear'),
                feats[1],
                feats[2],
                feats[3],
                F.interpolate(feats[4], size=feats[3].shape[-2:], mode='bilinear'))
//  [python]
def interpolate_vocoder_input(scale_factor, spec):
    """Interpolate spectrogram by the scale factor.
    It is mainly used to match the sampling rates of
    the tts and vocoder models.

    Args:
        scale_factor (float): scale factor to interpolate the spectrogram
        spec (np.array): spectrogram to be interpolated

    Returns:
        torch.tensor: interpolated spectrogram.
    """
    print(" > before interpolation :", spec.shape)
    spec = torch.tensor(spec).unsqueeze(0).unsqueeze(0)  # pylint: disable=not-callable
    spec = torch.nn.functional.interpolate(spec,
                                           scale_factor=scale_factor,
                                           recompute_scale_factor=True,
                                           mode='bilinear',
                                           align_corners=False).squeeze(0)
    print(" > after interpolation :", spec.shape)
    return spec
//  [python]
def resize(clip, target_size, interpolation_mode):
    if len(target_size) != 2:
        raise ValueError(
            f"target size should be tuple (height, width), instead got {target_size}"
        )
    return torch.nn.functional.interpolate(
        clip, size=target_size, mode=interpolation_mode, align_corners=False
    )
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`DPTForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If left to
                None, predictions will not be resized.
        Returns:
            semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = t.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`DPTForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If left to
                None, predictions will not be resized.
        Returns:
            semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = torch.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation
//  [python]
def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"):
        log = dict()
        loss = (latent_inputs - latent_predictions) ** 2
        log[f"{split}/latent_l2_loss"] = loss.mean().detach()
        image_reconstructions = None
        if self.perceptual_weight > 0.0:
            image_reconstructions = self.decoder.decode(latent_predictions)
            image_targets = self.decoder.decode(latent_inputs)
            perceptual_loss = self.perceptual_loss(
                image_targets.contiguous(), image_reconstructions.contiguous()
            )
            loss = (
                self.latent_weight * loss.mean()
                + self.perceptual_weight * perceptual_loss.mean()
            )
            log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach()

        if self.perceptual_weight_on_inputs > 0.0:
            image_reconstructions = default(
                image_reconstructions, self.decoder.decode(latent_predictions)
            )
            if self.scale_input_to_tgt_size:
                image_inputs = t.nn.functional.interpolate(
                    image_inputs,
                    image_reconstructions.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )
            elif self.scale_tgt_to_input_size:
                image_reconstructions = t.nn.functional.interpolate(
                    image_reconstructions,
                    image_inputs.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )

            perceptual_loss2 = self.perceptual_loss(
                image_inputs.contiguous(), image_reconstructions.contiguous()
            )
            loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean()
            log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach()
        return loss, log
//  [python]
def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"):
        log = dict()
        loss = (latent_inputs - latent_predictions) ** 2
        log[f"{split}/latent_l2_loss"] = loss.mean().detach()
        image_reconstructions = None
        if self.perceptual_weight > 0.0:
            image_reconstructions = self.decoder.decode(latent_predictions)
            image_targets = self.decoder.decode(latent_inputs)
            perceptual_loss = self.perceptual_loss(
                image_targets.contiguous(), image_reconstructions.contiguous()
            )
            loss = (
                self.latent_weight * loss.mean()
                + self.perceptual_weight * perceptual_loss.mean()
            )
            log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach()

        if self.perceptual_weight_on_inputs > 0.0:
            image_reconstructions = default(
                image_reconstructions, self.decoder.decode(latent_predictions)
            )
            if self.scale_input_to_tgt_size:
                image_inputs = torch.nn.functional.interpolate(
                    image_inputs,
                    image_reconstructions.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )
            elif self.scale_tgt_to_input_size:
                image_reconstructions = torch.nn.functional.interpolate(
                    image_reconstructions,
                    image_inputs.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )

            perceptual_loss2 = self.perceptual_loss(
                image_inputs.contiguous(), image_reconstructions.contiguous()
            )
            loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean()
            log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach()
        return loss, log
//  [python]
def forward(self, x1, x2):
        batch_size = x1.shape[0]
        if self.downsample is not None:
            x1, x2 = F.interpolate(x1, self.downsample), F.interpolate(x2, self.downsample)
        self.features_extractor(torch.cat([x1, x2], dim=1))
        features = self.features.output.view([batch_size, -1])
        logits = self.type_estimator(features)
        shift = self.shift_estimator(features)
        return logits, shift.squeeze()
//  [python]
def view_alphas_to_tensor(
        self, view_alphas: Union[torch.Tensor, List[List[Image.Image]]]
    ) -> torch.Tensor:
        """
        Returns a [batch x num_views x 1 x size x size] tensor in the range [0, 1].
        """
        if isinstance(view_alphas, torch.Tensor):
            return view_alphas

        tensor_batch = []
        num_views = len(view_alphas[0])
        for inner_list in view_alphas:
            assert len(inner_list) == num_views
            inner_batch = []
            for img in inner_list:
                tensor = (
                    torch.from_numpy(np.array(img)).to(device=self.device, dtype=torch.float32)
                    / 255.0
                )
                tensor = t.interpolate(
                    tensor[None, None],
                    (self.image_size,) * 2,
                    mode="nearest",
                )
                inner_batch.append(tensor)
            tensor_batch.append(torch.cat(inner_batch, dim=0))
        return torch.stack(tensor_batch, dim=0)
//  [python]
def view_alphas_to_tensor(
        self, view_alphas: Union[torch.Tensor, List[List[Image.Image]]]
    ) -> torch.Tensor:
        """
        Returns a [batch x num_views x 1 x size x size] tensor in the range [0, 1].
        """
        if isinstance(view_alphas, torch.Tensor):
            return view_alphas

        tensor_batch = []
        num_views = len(view_alphas[0])
        for inner_list in view_alphas:
            assert len(inner_list) == num_views
            inner_batch = []
            for img in inner_list:
                tensor = (
                    torch.from_numpy(np.array(img)).to(device=self.device, dtype=torch.float32)
                    / 255.0
                )
                tensor = F.interpolate(
                    tensor[None, None],
                    (self.image_size,) * 2,
                    mode="nearest",
                )
                inner_batch.append(tensor)
            tensor_batch.append(torch.cat(inner_batch, dim=0))
        return torch.stack(tensor_batch, dim=0)
//  [python]
def visualize_results_vss(self, batched_inputs, sem_masks_dict, out_size, sub_task):
        file_names = batched_inputs[0]['file_names'] 
        dataset_name = file_names[0].split('/')[1]
        video_id = file_names[0].split('/')[-2]

        save_dir = os.path.join(self.output_dir, 'inference/vss_'+dataset_name, video_id)
        os.makedirs(save_dir, exist_ok=True)

        sem_masks = sem_masks_dict["pred_masks"]
        for t, sem_mask in enumerate(sem_masks):
            file_name = file_names[t].split('/')[-1]
            img = batched_inputs[0]["image"][t]
            img = t.interpolate(
                img[None].float(),
                out_size,
                mode="bilinear",
                align_corners=False
            ).squeeze(0).long()
            img = img.permute(1, 2, 0).cpu().to(torch.uint8)
            visualizer = VisualizerFrame(
                img, 
                metadata=self.metadata, 
                class_names=None
            )

            save_path = '/'.join([save_dir, file_name.replace('.png', '.ipg')])
            VisImage = visualizer.draw_sem_seg(sem_mask, alpha=0.6)
            VisImage.save(save_path)
        
        # save all frames to a .avi video
        out = cv2.VideoWriter(
            '/'.join([save_dir, video_id + '.avi']),
            cv2.VideoWriter_fourcc(*'DIVX'),
            fps=10,
            frameSize=(out_size[1], out_size[0])
        )

        file_names = glob.glob('/'.join([save_dir, "*.jpg"]))
        file_names = sorted(file_names, key=lambda f: int(f.split("/")[-1].split('_')[-1].replace('.jpg', '')))
        for file_name in file_names:
            out.write(cv2.imread(file_name))
        out.release()
        print(f"save all frames with sem. seg. into {video_id}.avi")
//  [python]
def visualize_results_vss(self, batched_inputs, sem_masks_dict, out_size, sub_task):
        file_names = batched_inputs[0]['file_names'] 
        dataset_name = file_names[0].split('/')[1]
        video_id = file_names[0].split('/')[-2]

        save_dir = os.path.join(self.output_dir, 'inference/vss_'+dataset_name, video_id)
        os.makedirs(save_dir, exist_ok=True)

        sem_masks = sem_masks_dict["pred_masks"]
        for t, sem_mask in enumerate(sem_masks):
            file_name = file_names[t].split('/')[-1]
            img = batched_inputs[0]["image"][t]
            img = F.interpolate(
                img[None].float(),
                out_size,
                mode="bilinear",
                align_corners=False
            ).squeeze(0).long()
            img = img.permute(1, 2, 0).cpu().to(torch.uint8)
            visualizer = VisualizerFrame(
                img, 
                metadata=self.metadata, 
                class_names=None
            )

            save_path = '/'.join([save_dir, file_name.replace('.png', '.ipg')])
            VisImage = visualizer.draw_sem_seg(sem_mask, alpha=0.6)
            VisImage.save(save_path)
        
        # save all frames to a .avi video
        out = cv2.VideoWriter(
            '/'.join([save_dir, video_id + '.avi']),
            cv2.VideoWriter_fourcc(*'DIVX'),
            fps=10,
            frameSize=(out_size[1], out_size[0])
        )

        file_names = glob.glob('/'.join([save_dir, "*.jpg"]))
        file_names = sorted(file_names, key=lambda f: int(f.split("/")[-1].split('_')[-1].replace('.jpg', '')))
        for file_name in file_names:
            out.write(cv2.imread(file_name))
        out.release()
        print(f"save all frames with sem. seg. into {video_id}.avi")
//  [python]
def raw_depths_to_tensor(
        self, depths: Union[torch.Tensor, List[List[Image.Image]]]
    ) -> torch.Tensor:
        """
        Returns a [batch x num_views x 1 x size x size] tensor
        """
        if isinstance(depths, torch.Tensor):
            return depths

        tensor_batch = []
        num_views = len(depths[0])
        for inner_list in depths:
            assert len(inner_list) == num_views
            inner_batch = []
            for arr in inner_list:
                tensor = torch.from_numpy(arr).clamp(max=self.max_depth)
                tensor = t.interpolate(
                    tensor[None, None],
                    (self.image_size,) * 2,
                    mode="nearest",
                )
                inner_batch.append(tensor.to(device=self.device, dtype=torch.float32))
            tensor_batch.append(torch.cat(inner_batch, dim=0))
        return torch.stack(tensor_batch, dim=0)
//  [python]
def raw_depths_to_tensor(
        self, depths: Union[torch.Tensor, List[List[Image.Image]]]
    ) -> torch.Tensor:
        """
        Returns a [batch x num_views x 1 x size x size] tensor
        """
        if isinstance(depths, torch.Tensor):
            return depths

        tensor_batch = []
        num_views = len(depths[0])
        for inner_list in depths:
            assert len(inner_list) == num_views
            inner_batch = []
            for arr in inner_list:
                tensor = torch.from_numpy(arr).clamp(max=self.max_depth)
                tensor = F.interpolate(
                    tensor[None, None],
                    (self.image_size,) * 2,
                    mode="nearest",
                )
                inner_batch.append(tensor.to(device=self.device, dtype=torch.float32))
            tensor_batch.append(torch.cat(inner_batch, dim=0))
        return torch.stack(tensor_batch, dim=0)