jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x: torch.Tensor) -> Union[torch.Tensor, List[torch.Tensor]]:
        _, _, h, w = x.shape

        # collect encode outputs
        encode_outputs = []
        for i, m in enumerate(self.encode_modules):
            x = m(x)
            encode_outputs.append(x)
            if i != self.encode_num - 1:
                x = F.max_pool2d(x, kernel_size=2, stride=2, ceil_mode=True)

        # collect decode outputs
        x = encode_outputs.pop()
        decode_outputs = [x]
        for m in self.decode_modules:
            x2 = encode_outputs.pop()
            x = F.interpolate(x, size=x2.shape[2:], mode='bilinear', align_corners=False)
            x = m(torch.concat([x, x2], dim=1))
            decode_outputs.insert(0, x)

        # collect side outputs
        side_outputs = []
        for m in self.side_modules:
            x = decode_outputs.pop()
            x = F.interpolate(m(x), size=[h, w], mode='bilinear', align_corners=False)
            side_outputs.insert(0, x)

        x = self.out_conv(torch.concat(side_outputs, dim=1))

        if self.training:
            # do not use torch.sigmoid for amp safe
            return [x] + side_outputs
        else:
            return torch.sigmoid(x)
//  [python]
def interpolate_1d(x, length, mode = 'bilinear'):
    x = rearrange(x, 'b c t -> b c t 1')
    x = t.interpolate(x, (length, 1), mode = mode)
    return rearrange(x, 'b c t 1 -> b c t')
//  [python]
def interpolate_1d(x, length, mode = 'bilinear'):
    x = rearrange(x, 'b c t -> b c t 1')
    x = F.interpolate(x, (length, 1), mode = mode)
    return rearrange(x, 'b c t 1 -> b c t')
//  [python]
def get_rgbd_point_cloud(camera: CamerasBase, image_rgb: torch.Tensor, depth_map: torch.Tensor, mask: Optional[torch.Tensor]=None, mask_thr: float=0.5, mask_points: bool=True) ->Pointclouds:
    """
    Given a batch of images, depths, masks and cameras, generate a colored
    point cloud by unprojecting depth maps to the  and coloring with the source
    pixel colors.
    """
    imh, imw = image_rgb.shape[2:]
    pts_3d = ray_bundle_to_ray_points(NDCMultinomialRaysampler(image_width=imw, image_height=imh, n_pts_per_ray=1, min_depth=1.0, max_depth=1.0)(camera)._replace(lengths=depth_map[:, 0, ..., None]))
    pts_mask = depth_map > 0.0
    if mask is not None:
        pts_mask *= mask > mask_thr
    pts_mask = pts_mask.reshape(-1)
    pts_3d = pts_3d.reshape(-1, 3)[pts_mask]
    pts_colors = t.nn.functional.interpolate(image_rgb, size=[imh, imw], mode='bilinear', align_corners=False)
    pts_colors = pts_colors.permute(0, 2, 3, 1).reshape(-1, 3)[pts_mask]
    return Pointclouds(points=pts_3d[None], features=pts_colors[None])
//  [python]
def get_rgbd_point_cloud(camera: CamerasBase, image_rgb: torch.Tensor, depth_map: torch.Tensor, mask: Optional[torch.Tensor]=None, mask_thr: float=0.5, mask_points: bool=True) ->Pointclouds:
    """
    Given a batch of images, depths, masks and cameras, generate a colored
    point cloud by unprojecting depth maps to the  and coloring with the source
    pixel colors.
    """
    imh, imw = image_rgb.shape[2:]
    pts_3d = ray_bundle_to_ray_points(NDCMultinomialRaysampler(image_width=imw, image_height=imh, n_pts_per_ray=1, min_depth=1.0, max_depth=1.0)(camera)._replace(lengths=depth_map[:, 0, ..., None]))
    pts_mask = depth_map > 0.0
    if mask is not None:
        pts_mask *= mask > mask_thr
    pts_mask = pts_mask.reshape(-1)
    pts_3d = pts_3d.reshape(-1, 3)[pts_mask]
    pts_colors = torch.nn.functional.interpolate(image_rgb, size=[imh, imw], mode='bilinear', align_corners=False)
    pts_colors = pts_colors.permute(0, 2, 3, 1).reshape(-1, 3)[pts_mask]
    return Pointclouds(points=pts_3d[None], features=pts_colors[None])
//  [python]
def downsample(mask: torch.FloatTensor, batch_size: int, num_queries: int, value_embed_dim: int):
        """
        Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention. If the
        aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.

        Args:
            mask (`torch.FloatTensor`):
                The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
            batch_size (`int`):
                The batch size.
            num_queries (`int`):
                The number of queries.
            value_embed_dim (`int`):
                The dimensionality of the value embeddings.

        Returns:
            `torch.FloatTensor`:
                The downsampled mask tensor.

        """
        o_h = mask.shape[1]
        o_w = mask.shape[2]
        ratio = o_w / o_h
        mask_h = int(math.sqrt(num_queries / ratio))
        mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
        mask_w = num_queries // mask_h

        mask_downsample = t.nn.functional.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)

        # Repeat batch_size times
        if mask_downsample.shape[0] < batch_size:
            mask_downsample = mask_downsample.repeat(batch_size, 1, 1)

        mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)

        downsampled_area = mask_h * mask_w
        # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
        # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
        if downsampled_area < num_queries:
            warnings.warn(
                "The aspect ratio of the mask does not match the aspect ratio of the output image. "
                "Please update your masks or adjust the output size for optimal performance.",
                UserWarning,
            )
            mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
        # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
        if downsampled_area > num_queries:
            warnings.warn(
                "The aspect ratio of the mask does not match the aspect ratio of the output image. "
                "Please update your masks or adjust the output size for optimal performance.",
                UserWarning,
            )
            mask_downsample = mask_downsample[:, :num_queries]

        # Repeat last dimension to match SDPA output shape
        mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
            1, 1, value_embed_dim
        )

        return mask_downsample
//  [python]
def downsample(mask: torch.FloatTensor, batch_size: int, num_queries: int, value_embed_dim: int):
        """
        Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention. If the
        aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.

        Args:
            mask (`torch.FloatTensor`):
                The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
            batch_size (`int`):
                The batch size.
            num_queries (`int`):
                The number of queries.
            value_embed_dim (`int`):
                The dimensionality of the value embeddings.

        Returns:
            `torch.FloatTensor`:
                The downsampled mask tensor.

        """
        o_h = mask.shape[1]
        o_w = mask.shape[2]
        ratio = o_w / o_h
        mask_h = int(math.sqrt(num_queries / ratio))
        mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
        mask_w = num_queries // mask_h

        mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)

        # Repeat batch_size times
        if mask_downsample.shape[0] < batch_size:
            mask_downsample = mask_downsample.repeat(batch_size, 1, 1)

        mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)

        downsampled_area = mask_h * mask_w
        # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
        # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
        if downsampled_area < num_queries:
            warnings.warn(
                "The aspect ratio of the mask does not match the aspect ratio of the output image. "
                "Please update your masks or adjust the output size for optimal performance.",
                UserWarning,
            )
            mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
        # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
        if downsampled_area > num_queries:
            warnings.warn(
                "The aspect ratio of the mask does not match the aspect ratio of the output image. "
                "Please update your masks or adjust the output size for optimal performance.",
                UserWarning,
            )
            mask_downsample = mask_downsample[:, :num_queries]

        # Repeat last dimension to match SDPA output shape
        mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
            1, 1, value_embed_dim
        )

        return mask_downsample
//  [python]
def 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 in the range [-1, 1].
        """
        if isinstance(depths, torch.Tensor):
            return depths

        tensor_batch = []
        for inner_list in depths:
            assert len(inner_list) == self.num_views
            inner_batch = []
            for arr in inner_list:
                tensor = torch.from_numpy(arr).clamp(max=self.max_depth) / self.max_depth
                tensor = tensor * 2 - 1
                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 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 in the range [-1, 1].
        """
        if isinstance(depths, torch.Tensor):
            return depths

        tensor_batch = []
        for inner_list in depths:
            assert len(inner_list) == self.num_views
            inner_batch = []
            for arr in inner_list:
                tensor = torch.from_numpy(arr).clamp(max=self.max_depth) / self.max_depth
                tensor = tensor * 2 - 1
                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)
//  [python]
def forward(self, imgs: Optional[torch.Tensor], masks: Optional[torch.Tensor]=None, **kwargs) ->Dict[Any, torch.Tensor]:
        """
        Args:
            imgs: A batch of input images of shape `(B, 3, H, W)`.
            masks: A batch of input masks of shape `(B, 3, H, W)`.

        Returns:
            out_feats: A dict `{f_i: t_i}` keyed by predicted feature names `f_i`
                and their corresponding tensors `t_i` of shape `(B, dim_i, H_i, W_i)`.
        """
        out_feats = {}
        imgs_input = imgs
        if self.image_rescale != 1.0 and imgs_input is not None:
            imgs_resized = t.interpolate(imgs_input, scale_factor=self.image_rescale, mode='bilinear')
        else:
            imgs_resized = imgs_input
        if len(self.stages) > 0:
            assert imgs_resized is not None
            if self.normalize_image:
                imgs_normed = self._resnet_normalize_image(imgs_resized)
            else:
                imgs_normed = imgs_resized
            feats = self.stem(imgs_normed)
            for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
                feats = layer(feats)
                assert feats.shape[1] == _FEAT_DIMS[self.name][stage]
                if stage + 1 in self.stages:
                    f = proj(feats)
                    if self.global_average_pool:
                        f = f.mean(dims=(2, 3))
                    if self.l2_norm:
                        normfac = 1.0 / math.sqrt(len(self.stages))
                        f = Fu.normalize(f, dim=1) * normfac
                    feature_name = self._get_resnet_stage_feature_name(stage)
                    out_feats[feature_name] = f
        if self.add_masks:
            assert masks is not None
            out_feats[MASK_FEATURE_NAME] = masks
        if self.add_images:
            assert imgs_resized is not None
            out_feats[IMAGE_FEATURE_NAME] = imgs_resized
        if self.feature_rescale != 1.0:
            out_feats = {k: (self.feature_rescale * f) for k, f in out_feats.items()}
        return out_feats
//  [python]
def forward(self, imgs: Optional[torch.Tensor], masks: Optional[torch.Tensor]=None, **kwargs) ->Dict[Any, torch.Tensor]:
        """
        Args:
            imgs: A batch of input images of shape `(B, 3, H, W)`.
            masks: A batch of input masks of shape `(B, 3, H, W)`.

        Returns:
            out_feats: A dict `{f_i: t_i}` keyed by predicted feature names `f_i`
                and their corresponding tensors `t_i` of shape `(B, dim_i, H_i, W_i)`.
        """
        out_feats = {}
        imgs_input = imgs
        if self.image_rescale != 1.0 and imgs_input is not None:
            imgs_resized = Fu.interpolate(imgs_input, scale_factor=self.image_rescale, mode='bilinear')
        else:
            imgs_resized = imgs_input
        if len(self.stages) > 0:
            assert imgs_resized is not None
            if self.normalize_image:
                imgs_normed = self._resnet_normalize_image(imgs_resized)
            else:
                imgs_normed = imgs_resized
            feats = self.stem(imgs_normed)
            for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
                feats = layer(feats)
                assert feats.shape[1] == _FEAT_DIMS[self.name][stage]
                if stage + 1 in self.stages:
                    f = proj(feats)
                    if self.global_average_pool:
                        f = f.mean(dims=(2, 3))
                    if self.l2_norm:
                        normfac = 1.0 / math.sqrt(len(self.stages))
                        f = Fu.normalize(f, dim=1) * normfac
                    feature_name = self._get_resnet_stage_feature_name(stage)
                    out_feats[feature_name] = f
        if self.add_masks:
            assert masks is not None
            out_feats[MASK_FEATURE_NAME] = masks
        if self.add_images:
            assert imgs_resized is not None
            out_feats[IMAGE_FEATURE_NAME] = imgs_resized
        if self.feature_rescale != 1.0:
            out_feats = {k: (self.feature_rescale * f) for k, f in out_feats.items()}
        return out_feats
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        """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 resize_and_pad(self, feauture_maps, info_array):
        """Resize the instance feature as well as the weight map to match the
        size of full-image and do zero padding on both of them.

        Args:
            feauture_maps (tensor): Feature map
            info_array (tensor): The bounding box

        Returns:
            feauture_maps (tensor): Feature maps after resize and padding
        """
        feauture_maps = torch.nn.functional.interpolate(
            feauture_maps,
            size=(info_array[5], info_array[4]),
            mode='bilinear')
        feauture_maps = torch.nn.functional.pad(feauture_maps,
                                                (info_array[0], info_array[1],
                                                 info_array[2], info_array[3]),
                                                'constant', 0)
        return feauture_maps
//  [python]
def _resize_image_and_masks(image: Tensor,
                            self_min_size: float,
                            self_max_size: float,
                            target: Optional[Dict[str, Tensor]] = None,
                            fixed_size: Optional[Tuple[int, int]] = None
                            ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:

    if torchvision._is_tracing():
        im_shape = _get_shape_onnx(image)
    else:
        im_shape = torch.tensor(image.shape[-2:])

    size: Optional[List[int]] = None
    scale_factor: Optional[float] = None
    recompute_scale_factor: Optional[bool] = None
    if fixed_size is not None:
        size = [fixed_size[1], fixed_size[0]]
    else:
        min_size = torch.min(im_shape).to(dtype=torch.float32)  # 获取高宽中的最小值
        max_size = torch.max(im_shape).to(dtype=torch.float32)  # 获取高宽中的最大值
        scale = torch.min(self_min_size / min_size, self_max_size / max_size)  # 计算缩放比例

        if torchvision._is_tracing():
            scale_factor = _fake_cast_onnx(scale)
        else:
            scale_factor = scale.item()
        recompute_scale_factor = True

    # interpolate利用插值的方法缩放图片
    # image[None]操作是在最前面添加batch维度[C, H, W] -> [1, C, H, W]
    # bilinear只支持4D Tensor
    image = t.interpolate(
        image[None],
        size=size,
        scale_factor=scale_factor,
        mode="bilinear",
        align_corners=False)[0]

    if target is None:
        return image, target

    if "masks" in target:
        mask = target["masks"]
        mask = t.interpolate(
            mask[:, None].float(), size=size, scale_factor=scale_factor,
        )[:, 0].byte()  # self.byte() is equivalent to self.to(torch.uint8).
        target["masks"] = mask

    return image, target
//  [python]
def _resize_image_and_masks(image: Tensor,
                            self_min_size: float,
                            self_max_size: float,
                            target: Optional[Dict[str, Tensor]] = None,
                            fixed_size: Optional[Tuple[int, int]] = None
                            ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:

    if torchvision._is_tracing():
        im_shape = _get_shape_onnx(image)
    else:
        im_shape = torch.tensor(image.shape[-2:])

    size: Optional[List[int]] = None
    scale_factor: Optional[float] = None
    recompute_scale_factor: Optional[bool] = None
    if fixed_size is not None:
        size = [fixed_size[1], fixed_size[0]]
    else:
        min_size = torch.min(im_shape).to(dtype=torch.float32)  # 获取高宽中的最小值
        max_size = torch.max(im_shape).to(dtype=torch.float32)  # 获取高宽中的最大值
        scale = torch.min(self_min_size / min_size, self_max_size / max_size)  # 计算缩放比例

        if torchvision._is_tracing():
            scale_factor = _fake_cast_onnx(scale)
        else:
            scale_factor = scale.item()
        recompute_scale_factor = True

    # interpolate利用插值的方法缩放图片
    # image[None]操作是在最前面添加batch维度[C, H, W] -> [1, C, H, W]
    # bilinear只支持4D Tensor
    image = torch.nn.functional.interpolate(
        image[None],
        size=size,
        scale_factor=scale_factor,
        mode="bilinear",
        recompute_scale_factor=recompute_scale_factor,
        align_corners=False)[0]

    if target is None:
        return image, target

    if "masks" in target:
        mask = target["masks"]
        mask = torch.nn.functional.interpolate(
            mask[:, None].float(), size=size, scale_factor=scale_factor, recompute_scale_factor=recompute_scale_factor
        )[:, 0].byte()  # self.byte() is equivalent to self.to(torch.uint8).
        target["masks"] = mask

    return image, target
//  [python]
def _onnx_paste_mask_in_image(mask, box, im_h, im_w):
    one = torch.ones(1, dtype=torch.int64)
    zero = torch.zeros(1, dtype=torch.int64)

    w = box[2] - box[0] + one
    h = box[3] - box[1] + one
    w = torch.max(torch.cat((w, one)))
    h = torch.max(torch.cat((h, one)))

    # Set shape to [batchxCxHxW]
    mask = mask.expand((1, 1, mask.size(0), mask.size(1)))

    # Resize mask
    mask = t.nn.functional.interpolate(mask, size=(int(h), int(w)), mode="bilinear", align_corners=False)
    mask = mask[0][0]

    x_0 = torch.max(torch.cat((box[0].unsqueeze(0), zero)))
    x_1 = torch.min(torch.cat((box[2].unsqueeze(0) + one, im_w.unsqueeze(0))))
    y_0 = torch.max(torch.cat((box[1].unsqueeze(0), zero)))
    y_1 = torch.min(torch.cat((box[3].unsqueeze(0) + one, im_h.unsqueeze(0))))

    unpaded_im_mask = mask[(y_0 - box[1]): (y_1 - box[1]), (x_0 - box[0]): (x_1 - box[0])]

    # TODO : replace below with a dynamic padding when support is added in ONNX

    # pad y
    zeros_y0 = torch.zeros(y_0, unpaded_im_mask.size(1))
    zeros_y1 = torch.zeros(im_h - y_1, unpaded_im_mask.size(1))
    concat_0 = torch.cat((zeros_y0, unpaded_im_mask.to(dtype=torch.float32), zeros_y1), 0)[0:im_h, :]
    # pad x
    zeros_x0 = torch.zeros(concat_0.size(0), x_0)
    zeros_x1 = torch.zeros(concat_0.size(0), im_w - x_1)
    im_mask = torch.cat((zeros_x0, concat_0, zeros_x1), 1)[:, :im_w]
    return im_mask
//  [python]
def _onnx_paste_mask_in_image(mask, box, im_h, im_w):
    one = torch.ones(1, dtype=torch.int64)
    zero = torch.zeros(1, dtype=torch.int64)

    w = box[2] - box[0] + one
    h = box[3] - box[1] + one
    w = torch.max(torch.cat((w, one)))
    h = torch.max(torch.cat((h, one)))

    # Set shape to [batchxCxHxW]
    mask = mask.expand((1, 1, mask.size(0), mask.size(1)))

    # Resize mask
    mask = F.interpolate(mask, size=(int(h), int(w)), mode="bilinear", align_corners=False)
    mask = mask[0][0]

    x_0 = torch.max(torch.cat((box[0].unsqueeze(0), zero)))
    x_1 = torch.min(torch.cat((box[2].unsqueeze(0) + one, im_w.unsqueeze(0))))
    y_0 = torch.max(torch.cat((box[1].unsqueeze(0), zero)))
    y_1 = torch.min(torch.cat((box[3].unsqueeze(0) + one, im_h.unsqueeze(0))))

    unpaded_im_mask = mask[(y_0 - box[1]): (y_1 - box[1]), (x_0 - box[0]): (x_1 - box[0])]

    # TODO : replace below with a dynamic padding when support is added in ONNX

    # pad y
    zeros_y0 = torch.zeros(y_0, unpaded_im_mask.size(1))
    zeros_y1 = torch.zeros(im_h - y_1, unpaded_im_mask.size(1))
    concat_0 = torch.cat((zeros_y0, unpaded_im_mask.to(dtype=torch.float32), zeros_y1), 0)[0:im_h, :]
    # pad x
    zeros_x0 = torch.zeros(concat_0.size(0), x_0)
    zeros_x1 = torch.zeros(concat_0.size(0), im_w - x_1)
    im_mask = torch.cat((zeros_x0, concat_0, zeros_x1), 1)[:, :im_w]
    return im_mask
//  [python]
def resize(image, size):
    image = t.imresize(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
    return image
//  [python]
def resize(image, size):
    image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
    return image
//  [python]
def paste_mask_in_image(mask, box, im_h, im_w):
    # type: (Tensor, Tensor, int, int) -> Tensor

    # refer to: https://github.com/pytorch/vision/issues/5845
    TO_REMOVE = 1
    w = int(box[2] - box[0] + TO_REMOVE)
    h = int(box[3] - box[1] + TO_REMOVE)
    w = max(w, 1)
    h = max(h, 1)

    # Set shape to [batch, C, H, W]
    # 因为后续的bilinear操作只支持4-D的Tensor
    mask = mask.expand((1, 1, -1, -1))  # -1 means not changing the size of that dimension

    # Resize mask
    mask = t.nn.functional.interpolate(mask, size=(h, w), mode='bilinear', align_corners=False)
    mask = mask[0][0]  # [batch, C, H, W] -> [H, W]

    im_mask = torch.zeros((im_h, im_w), dtype=mask.dtype, device=mask.device)
    # 填入原图的目标区域(防止越界)
    x_0 = max(box[0], 0)
    x_1 = min(box[2] + 1, im_w)
    y_0 = max(box[1], 0)
    y_1 = min(box[3] + 1, im_h)

    # 将resize后的mask填入对应目标区域
    im_mask[y_0:y_1, x_0:x_1] = mask[(y_0 - box[1]):(y_1 - box[1]), (x_0 - box[0]):(x_1 - box[0])]
    return im_mask
//  [python]
def paste_mask_in_image(mask, box, im_h, im_w):
    # type: (Tensor, Tensor, int, int) -> Tensor

    # refer to: https://github.com/pytorch/vision/issues/5845
    TO_REMOVE = 1
    w = int(box[2] - box[0] + TO_REMOVE)
    h = int(box[3] - box[1] + TO_REMOVE)
    w = max(w, 1)
    h = max(h, 1)

    # Set shape to [batch, C, H, W]
    # 因为后续的bilinear操作只支持4-D的Tensor
    mask = mask.expand((1, 1, -1, -1))  # -1 means not changing the size of that dimension

    # Resize mask
    mask = F.interpolate(mask, size=(h, w), mode='bilinear', align_corners=False)
    mask = mask[0][0]  # [batch, C, H, W] -> [H, W]

    im_mask = torch.zeros((im_h, im_w), dtype=mask.dtype, device=mask.device)
    # 填入原图的目标区域(防止越界)
    x_0 = max(box[0], 0)
    x_1 = min(box[2] + 1, im_w)
    y_0 = max(box[1], 0)
    y_1 = min(box[3] + 1, im_h)

    # 将resize后的mask填入对应目标区域
    im_mask[y_0:y_1, x_0:x_1] = mask[(y_0 - box[1]):(y_1 - box[1]), (x_0 - box[0]):(x_1 - box[0])]
    return im_mask
//  [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 forward(self, x):
        return self.conv(F.upsample(x, scale_factor=2, mode='nearest'))
//  [python]
def forward(self, x):
        return self.conv(F.interpolate(x, scale_factor=2, mode='nearest'))
//  [python]
def forward(self, input):
        # names = list(input.keys())
        # input = list(input.values())

        output1 = self.output1(input[0])
        output2 = self.output2(input[1])
        output3 = self.output3(input[2])

        up3 = F.upsample(output3, size=[output2.size(2), output2.size(3)], mode='nearest')
        output2 = output2 + up3
        output2 = self.merge2(output2)

        up2 = F.upsample(output2, size=[output1.size(2), output1.size(3)], mode='nearest')
        output1 = output1 + up2
        output1 = self.merge1(output1)

        out = [output1, output2, output3]
        return out
//  [python]
def forward(self, input):
        # names = list(input.keys())
        # input = list(input.values())

        output1 = self.output1(input[0])
        output2 = self.output2(input[1])
        output3 = self.output3(input[2])

        up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode='nearest')
        output2 = output2 + up3
        output2 = self.merge2(output2)

        up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode='nearest')
        output1 = output1 + up2
        output1 = self.merge1(output1)

        out = [output1, output2, output3]
        return out
//  [python]
def forward(self, z, img):
        b, c, h, w = img.shape

        feat = self.mlp(z).view((-1, 256, 2, 2))
        warp = self.upsample(feat) / self.tex_size
        grid = warp + self.normal_grid
        if grid.shape[2] < img.shape[2]:
            grid = F.upsample(grid, scale_factor=img.shape[2] / grid.shape[2])
        grid = grid.permute(0, 2, 3, 1).contiguous()

        out = F.grid_sample(img, grid)
        return out, grid
//  [python]
def forward(self, z, img):
        b, c, h, w = img.shape

        feat = self.mlp(z).view((-1, 256, 2, 2))
        warp = self.upsample(feat) / self.tex_size
        grid = warp + self.normal_grid
        if grid.shape[2] < img.shape[2]:
            grid = F.interpolate(grid, scale_factor=img.shape[2] / grid.shape[2])
        grid = grid.permute(0, 2, 3, 1).contiguous()

        out = F.grid_sample(img, grid)
        return out, grid
//  [python]
def forward(self, x):
        if self.use_bilinear:
            x = F.upsample(x, scale_factor=2, mode='bilinear')
        out = self.deconv(x) + self.bias
        return out
//  [python]
def forward(self, x):
        if self.use_bilinear:
            x = F.interpolate(x, scale_factor=2)
        out = self.deconv(x) + self.bias
        return out
//  [python]
def forward(self, z, mask):  ## The shape of mask is Batch*1*H*W

        x=self.res_block(z)

        b,c,h,w=x.shape

        ## mask resize + dilation
        # tmp = 1 - mask
        mask = F.upsample(mask, size=(x.size(2), x.size(3)), mode="bilinear")
        mask[mask > 0] = 1.0

        # mask = 1 - mask
        # tmp = F.interpolate(tmp, (x.size(2), x.size(3)))
        # mask *= tmp
        # mask=1-mask
        ## 1: mask position 0: non-mask

        mask_unfold=F.unfold(mask, kernel_size=(self.patch_size,self.patch_size), padding=0, stride=self.patch_size)
        non_mask_region=(torch.mean(mask_unfold,dim=1,keepdim=True)>0.6).float()
        all_patch_num=h*w/self.patch_size/self.patch_size
        non_mask_region=non_mask_region.repeat(1,int(all_patch_num),1)

        x_unfold=F.unfold(x, kernel_size=(self.patch_size,self.patch_size), padding=0, stride=self.patch_size)
        y_unfold=x_unfold.permute(0,2,1)
        x_unfold_normalized=F.normalize(x_unfold,dim=1)
        y_unfold_normalized=F.normalize(y_unfold,dim=2)
        correlation_matrix=torch.bmm(y_unfold_normalized,x_unfold_normalized)
        correlation_matrix=correlation_matrix.masked_fill(non_mask_region==1.,-1e9)
        correlation_matrix=F.softmax(correlation_matrix,dim=2)

        # print(correlation_matrix)

        R, max_arg=torch.max(correlation_matrix,dim=2)

        composed_unfold=self.Hard_Compose(x_unfold, 2, max_arg)
        composed_fold=F.fold(composed_unfold,output_size=(h,w),kernel_size=(self.patch_size,self.patch_size),padding=0,stride=self.patch_size)

        concat_1=torch.cat((z,composed_fold,mask),dim=1)
        concat_1=self.F_Combine(concat_1)

        return concat_1
//  [python]
def forward(self, z, mask):  ## The shape of mask is Batch*1*H*W

        x=self.res_block(z)

        b,c,h,w=x.shape

        ## mask resize + dilation
        # tmp = 1 - mask
        mask = F.interpolate(mask, (x.size(2), x.size(3)), mode="bilinear")
        mask[mask > 0] = 1.0

        # mask = 1 - mask
        # tmp = F.interpolate(tmp, (x.size(2), x.size(3)))
        # mask *= tmp
        # mask=1-mask
        ## 1: mask position 0: non-mask

        mask_unfold=F.unfold(mask, kernel_size=(self.patch_size,self.patch_size), padding=0, stride=self.patch_size)
        non_mask_region=(torch.mean(mask_unfold,dim=1,keepdim=True)>0.6).float()
        all_patch_num=h*w/self.patch_size/self.patch_size
        non_mask_region=non_mask_region.repeat(1,int(all_patch_num),1)

        x_unfold=F.unfold(x, kernel_size=(self.patch_size,self.patch_size), padding=0, stride=self.patch_size)
        y_unfold=x_unfold.permute(0,2,1)
        x_unfold_normalized=F.normalize(x_unfold,dim=1)
        y_unfold_normalized=F.normalize(y_unfold,dim=2)
        correlation_matrix=torch.bmm(y_unfold_normalized,x_unfold_normalized)
        correlation_matrix=correlation_matrix.masked_fill(non_mask_region==1.,-1e9)
        correlation_matrix=F.softmax(correlation_matrix,dim=2)

        # print(correlation_matrix)

        R, max_arg=torch.max(correlation_matrix,dim=2)

        composed_unfold=self.Hard_Compose(x_unfold, 2, max_arg)
        composed_fold=F.fold(composed_unfold,output_size=(h,w),kernel_size=(self.patch_size,self.patch_size),padding=0,stride=self.patch_size)

        concat_1=torch.cat((z,composed_fold,mask),dim=1)
        concat_1=self.F_Combine(concat_1)

        return concat_1
//  [python]
def upsample2x(self, x):
        _, _, h, w = x.shape
        x = t.nn.functional.interpolate(x, size=(h*2, w*2),
                          mode='bilinear', align_corners=True)
        return x
//  [python]
def upsample2x(self, x):
        _, _, h, w = x.shape
        x = F.interpolate(x, size=(h*2, w*2),
                          mode='bilinear', align_corners=True)
        return x
//  [python]
def forward(self, x, mask):  ## The shape of mask is Batch*1*H*W
        batch_size = x.size(0)

        g_x = self.g(x).view(batch_size, self.inter_channels, -1)

        g_x = g_x.permute(0, 2, 1)

        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)

        theta_x = theta_x.permute(0, 2, 1)

        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)

        if self.cosin:
            theta_x = F.normalize(theta_x, dim=2)
            phi_x = F.normalize(phi_x, dim=1)

        f = torch.matmul(theta_x, phi_x)

        f /= self.temperature

        f_div_C = F.softmax(f, dim=2)

        tmp = 1 - mask
        mask = t.nn.functional.upsample_bilinear(mask, size=(x.size(2), x.size(3)))
        mask[mask > 0] = 1.0
        mask = 1 - mask

        tmp = t.nn.functional.upsample_bilinear(tmp, size=(x.size(2), x.size(3)))
        mask *= tmp

        mask_expand = mask.view(batch_size, 1, -1)
        mask_expand = mask_expand.repeat(1, x.size(2) * x.size(3), 1)

        # mask = 1 - mask
        # mask=F.interpolate(mask,(x.size(2),x.size(3)))
        # mask_expand=mask.view(batch_size,1,-1)
        # mask_expand=mask_expand.repeat(1,x.size(2)*x.size(3),1)

        if self.use_self:
            mask_expand[:, range(x.size(2) * x.size(3)), range(x.size(2) * x.size(3))] = 1.0

        #    print(mask_expand.shape)
        #    print(f_div_C.shape)

        f_div_C = mask_expand * f_div_C
        if self.renorm:
            f_div_C = F.normalize(f_div_C, p=1, dim=2)

        ###########################

        y = torch.matmul(f_div_C, g_x)

        y = y.permute(0, 2, 1).contiguous()

        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)

        W_y = self.res_block(W_y)

        if self.mode == "combine":
            full_mask = mask.repeat(1, self.inter_channels, 1, 1)
            z = full_mask * x + (1 - full_mask) * W_y
        return z
//  [python]
def forward(self, x, mask):  ## The shape of mask is Batch*1*H*W
        batch_size = x.size(0)

        g_x = self.g(x).view(batch_size, self.inter_channels, -1)

        g_x = g_x.permute(0, 2, 1)

        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)

        theta_x = theta_x.permute(0, 2, 1)

        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)

        if self.cosin:
            theta_x = F.normalize(theta_x, dim=2)
            phi_x = F.normalize(phi_x, dim=1)

        f = torch.matmul(theta_x, phi_x)

        f /= self.temperature

        f_div_C = F.softmax(f, dim=2)

        tmp = 1 - mask
        mask = F.interpolate(mask, (x.size(2), x.size(3)), mode="bilinear")
        mask[mask > 0] = 1.0
        mask = 1 - mask

        tmp = F.interpolate(tmp, (x.size(2), x.size(3)))
        mask *= tmp

        mask_expand = mask.view(batch_size, 1, -1)
        mask_expand = mask_expand.repeat(1, x.size(2) * x.size(3), 1)

        # mask = 1 - mask
        # mask=F.interpolate(mask,(x.size(2),x.size(3)))
        # mask_expand=mask.view(batch_size,1,-1)
        # mask_expand=mask_expand.repeat(1,x.size(2)*x.size(3),1)

        if self.use_self:
            mask_expand[:, range(x.size(2) * x.size(3)), range(x.size(2) * x.size(3))] = 1.0

        #    print(mask_expand.shape)
        #    print(f_div_C.shape)

        f_div_C = mask_expand * f_div_C
        if self.renorm:
            f_div_C = F.normalize(f_div_C, p=1, dim=2)

        ###########################

        y = torch.matmul(f_div_C, g_x)

        y = y.permute(0, 2, 1).contiguous()

        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)

        W_y = self.res_block(W_y)

        if self.mode == "combine":
            full_mask = mask.repeat(1, self.inter_channels, 1, 1)
            z = full_mask * x + (1 - full_mask) * W_y
        return z
//  [python]
def upsample(
        self, size: int, mode: str, align_corners: bool, consolidate_errors: bool = True
    ) -> "HybridZonotope":
        new_head = t.interpolate(
            self.head, size=size, mode=mode, align_corners=align_corners
        )
        delta = 0
        assert mode in ["nearest", "linear", "bilinear", "trilinear"], "Upsample"

        if self.beta is not None:
            new_beta = t.interpolate(
                self.beta, size=size, mode=mode, align_corners=align_corners
            )
            delta = delta + new_beta
        else:
            new_beta = None

        if self.errors is not None:
            errors_resized = self.errors.view(-1, *self.head.shape[1:])
            new_errors = t.interpolate(
                errors_resized, size=size, mode=mode, align_corners=align_corners
            )
            new_errors = new_errors.view(-1, *new_head.shape)
            delta = delta + new_errors.abs().sum(0)
        else:
            new_errors = None

        if consolidate_errors:
            return HybridZonotope.construct_from_bounds(
                new_head - delta, new_head + delta, domain=self.domain
            )
        else:
            return HybridZonotope(new_head, new_beta, new_errors, self.domain)
//  [python]
def upsample(
        self, size: int, mode: str, align_corners: bool, consolidate_errors: bool = True
    ) -> "HybridZonotope":
        new_head = F.interpolate(
            self.head, size=size, mode=mode, align_corners=align_corners
        )
        delta = 0
        assert mode in ["nearest", "linear", "bilinear", "trilinear"], "Upsample"

        if self.beta is not None:
            new_beta = F.interpolate(
                self.beta, size=size, mode=mode, align_corners=align_corners
            )
            delta = delta + new_beta
        else:
            new_beta = None

        if self.errors is not None:
            errors_resized = self.errors.view(-1, *self.head.shape[1:])
            new_errors = F.interpolate(
                errors_resized, size=size, mode=mode, align_corners=align_corners
            )
            new_errors = new_errors.view(-1, *new_head.shape)
            delta = delta + new_errors.abs().sum(0)
        else:
            new_errors = None

        if consolidate_errors:
            return HybridZonotope.construct_from_bounds(
                new_head - delta, new_head + delta, domain=self.domain
            )
        else:
            return HybridZonotope(new_head, new_beta, new_errors, self.domain)
//  [python]
def inference_forward(self,z,mask): ## Reduce the extra memory cost


        x=self.res_block(z)

        b,c,h,w=x.shape

        ## mask resize + dilation
        # tmp = 1 - mask
        mask = F.interpolate(mask, (x.size(2), x.size(3)), mode="bilinear")
        mask[mask > 0] = 1.0
        # mask = 1 - mask
        # tmp = F.interpolate(tmp, (x.size(2), x.size(3)))
        # mask *= tmp
        # mask=1-mask
        ## 1: mask position 0: non-mask

        mask_unfold=F.unfold(mask, kernel_size=(self.patch_size,self.patch_size), padding=0, stride=self.patch_size)
        non_mask_region=(torch.mean(mask_unfold,dim=1,keepdim=True)>0.6).float()[0,0,:] # 1*1*all_patch_num

        all_patch_num=h*w/self.patch_size/self.patch_size

        mask_index=torch.nonzero(non_mask_region,as_tuple=True)[0]


        if len(mask_index)==0: ## No mask patch is selected, no attention is needed

            composed_fold=x

        else:

            unmask_index=torch.nonzero(non_mask_region!=1,as_tuple=True)[0]

            x_unfold=F.unfold(x, kernel_size=(self.patch_size,self.patch_size), padding=0, stride=self.patch_size)
            
            Query_Patch=torch.index_select(x_unfold,2,mask_index)
            Key_Patch=torch.index_select(x_unfold,2,unmask_index)

            Query_Patch=Query_Patch.permute(0,2,1)        
            Query_Patch_normalized=F.normalize(Query_Patch,dim=2)
            Key_Patch_normalized=F.normalize(Key_Patch,dim=1)

            correlation_matrix=torch.bmm(Query_Patch_normalized,Key_Patch_normalized)
            correlation_matrix=F.softmax(correlation_matrix,dim=2)


            R, max_arg=torch.max(correlation_matrix,dim=2)

            composed_unfold=self.Hard_Compose(Key_Patch, 2, max_arg)
            x_unfold[:,:,mask_index]=composed_unfold
            composed_fold=F.fold(x_unfold,output_size=(h,w),kernel_size=(self.patch_size,self.patch_size),padding=0,stride=self.patch_size)

        concat_1=torch.cat((z,composed_fold,mask),dim=1)
        concat_1=self.F_Combine(concat_1)


        return concat_1
//  [python]
def _upsample_like(src,tar):

    # src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
    src = F.interpolate(src,size=tar.shape[2:],mode='bilinear', align_corners=True)

    return src
//  [python]
def forward(self, x: th.Tensor) -> th.Tensor:
        assert x.shape[1] == self.channels

        if self.dims == 3:
            t_factor = 1 if not self.third_up else self.scale_factor
            x = t.interpolate(
                x,
                (
                    t_factor * x.shape[2],
                    x.shape[3] * self.scale_factor,
                    x.shape[4] * self.scale_factor,
                ),
                mode="nearest",
            )
        else:
            x = t.interpolate(x, scale_factor=self.scale_factor, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x: th.Tensor) -> th.Tensor:
        assert x.shape[1] == self.channels

        if self.dims == 3:
            t_factor = 1 if not self.third_up else self.scale_factor
            x = F.interpolate(
                x,
                (
                    t_factor * x.shape[2],
                    x.shape[3] * self.scale_factor,
                    x.shape[4] * self.scale_factor,
                ),
                mode="nearest",
            )
        else:
            x = F.interpolate(x, scale_factor=self.scale_factor, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x, temb,light_emb):
        _, _, H, W = x.shape
        x = F.upsample(
            x, scale_factor=2, mode='nearest')
        x = self.main(x)
        return x
//  [python]
def forward(self, x, temb,light_emb):
        _, _, H, W = x.shape
        x = F.interpolate(
            x, scale_factor=2, mode='nearest')
        x = self.main(x)
        return x
//  [python]
def forward(self, hidden_state, residual=None, size=None):
        if residual is not None:
            if hidden_state.shape != residual.shape:
                residual = nn.functional.upsample(
                    residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear"
                )
            hidden_state = hidden_state + self.residual_layer1(residual)

        hidden_state = self.residual_layer2(hidden_state)

        modifier = {"scale_factor": 2} if size is None else {"size": size}

        hidden_state = nn.functional.upsample(
            hidden_state,
            **modifier,
            mode="bilinear",
        )
        hidden_state = self.projection(hidden_state)

        return hidden_state
//  [python]
def forward(self, hidden_state, residual=None, size=None):
        if residual is not None:
            if hidden_state.shape != residual.shape:
                residual = nn.functional.interpolate(
                    residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear", align_corners=False
                )
            hidden_state = hidden_state + self.residual_layer1(residual)

        hidden_state = self.residual_layer2(hidden_state)

        modifier = {"scale_factor": 2} if size is None else {"size": size}

        hidden_state = nn.functional.interpolate(
            hidden_state,
            **modifier,
            mode="bilinear",
            align_corners=True,
        )
        hidden_state = self.projection(hidden_state)

        return hidden_state
//  [python]
def _get_pos_embed(self, pos_embed, bcthw):

        if len(bcthw) == 4:
            t, h, w = 1, bcthw[-2], bcthw[-1]
        else:
            t, h, w = bcthw[-3], bcthw[-2], bcthw[-1]
        if self.cls_embed_on:
            cls_pos_embed = pos_embed[:, 0:1, :]
            pos_embed = pos_embed[:, 1:]
        txy_num = pos_embed.shape[1]
        p_t, p_h, p_w = self.patch_dims
        assert p_t * p_h * p_w == txy_num

        if (p_t, p_h, p_w) != (t, h, w):
            new_pos_embed = t.interpolate(
                pos_embed[:, :, :]
                .reshape(1, p_t, p_h, p_w, -1)
                .permute(0, 4, 1, 2, 3),
                size=(t, h, w),
                mode="trilinear",
            )
            pos_embed = new_pos_embed.reshape(1, -1, t * h * w).permute(0, 2, 1)

        if self.cls_embed_on:
            pos_embed = torch.cat((cls_pos_embed, pos_embed), dim=1)

        return pos_embed
//  [python]
def _get_pos_embed(self, pos_embed, bcthw):

        if len(bcthw) == 4:
            t, h, w = 1, bcthw[-2], bcthw[-1]
        else:
            t, h, w = bcthw[-3], bcthw[-2], bcthw[-1]
        if self.cls_embed_on:
            cls_pos_embed = pos_embed[:, 0:1, :]
            pos_embed = pos_embed[:, 1:]
        txy_num = pos_embed.shape[1]
        p_t, p_h, p_w = self.patch_dims
        assert p_t * p_h * p_w == txy_num

        if (p_t, p_h, p_w) != (t, h, w):
            new_pos_embed = F.interpolate(
                pos_embed[:, :, :]
                .reshape(1, p_t, p_h, p_w, -1)
                .permute(0, 4, 1, 2, 3),
                size=(t, h, w),
                mode="trilinear",
            )
            pos_embed = new_pos_embed.reshape(1, -1, t * h * w).permute(0, 2, 1)

        if self.cls_embed_on:
            pos_embed = torch.cat((cls_pos_embed, pos_embed), dim=1)

        return pos_embed
//  [python]
def forward(self, x):
        c2, c3, c4, c5 = x

        in5 = self.in5_conv(c5)
        in4 = self.in4_conv(c4)
        in3 = self.in3_conv(c3)
        in2 = self.in2_conv(c2)

        out4 = in4 + t.interpolate(
            in5,
            scale_factor=2,
            mode='nearest',
        )
        out3 = in3 + t.interpolate(
            out4,
            scale_factor=2,
            mode='nearest',
        )
        out2 = in2 + t.interpolate(
            out3,
            scale_factor=2,
            mode='nearest',
        )

        p5 = self.p5_conv(in5)
        p4 = self.p4_conv(out4)
        p3 = self.p3_conv(out3)
        p2 = self.p2_conv(out2)
        p5 = t.interpolate(
            p5,
            scale_factor=8,
            mode='nearest',
        )
        p4 = t.interpolate(
            p4,
            scale_factor=4,
            mode='nearest',
        )
        p3 = t.interpolate(
            p3,
            scale_factor=2,
            mode='nearest',
        )

        fuse = torch.cat([p5, p4, p3, p2], dim=1)
        return fuse