jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x, y, y_seg, registration=False):

        source, target = x, y
        x_in = torch.cat((x, y), 1)
        e0 = self.eninput(x_in)
        e0 = self.ec1(e0)

        e1 = self.ec2(e0)
        e1 = self.ec3(e1)

        e2 = self.ec4(e1)
        e2 = self.ec5(e2)

        e3 = self.ec6(e2)
        e3 = self.ec7(e3)

        e4 = self.ec8(e3)
        e4 = self.ec9(e4)

        d0 = torch.cat((self.up1(e4), e3), 1)

        d0 = self.dc1(d0)
        d0 = self.dc2(d0)

        d1 = torch.cat((self.up2(d0), e2), 1)

        d1 = self.dc3(d1)
        d1 = self.dc4(d1)

        d2 = torch.cat((self.up3(d1), e1), 1)

        d2 = self.dc5(d2)
        d2 = self.dc6(d2)

        flow_field = self.flow(d2, y_seg)

        if self.diff_int==1:
            preint_flow = flow_field / 2.
            flow_field = self.integrate(preint_flow) * 2.

        flow_field = t.upsample(flow_field, size=[160, 192, 224], mode='trilinear')
        y_source = self.transformer(source, flow_field)
        if self.diff_int==0:
            preint_flow = flow_field

        if not registration:
            return (y_source, preint_flow), flow_field
        else:
            return y_source, flow_field
//  [python]
def forward(self, x, y, y_seg, registration=False):

        source, target = x, y
        x_in = torch.cat((x, y), 1)
        e0 = self.eninput(x_in)
        e0 = self.ec1(e0)

        e1 = self.ec2(e0)
        e1 = self.ec3(e1)

        e2 = self.ec4(e1)
        e2 = self.ec5(e2)

        e3 = self.ec6(e2)
        e3 = self.ec7(e3)

        e4 = self.ec8(e3)
        e4 = self.ec9(e4)

        d0 = torch.cat((self.up1(e4), e3), 1)

        d0 = self.dc1(d0)
        d0 = self.dc2(d0)

        d1 = torch.cat((self.up2(d0), e2), 1)

        d1 = self.dc3(d1)
        d1 = self.dc4(d1)

        d2 = torch.cat((self.up3(d1), e1), 1)

        d2 = self.dc5(d2)
        d2 = self.dc6(d2)

        flow_field = self.flow(d2, y_seg)

        if self.diff_int==1:
            preint_flow = flow_field / 2.
            flow_field = self.integrate(preint_flow) * 2.

        flow_field = torch.nn.functional.interpolate(flow_field, size=[160, 192, 224], mode='trilinear', align_corners=False)
        y_source = self.transformer(source, flow_field)
        if self.diff_int==0:
            preint_flow = flow_field

        if not registration:
            return (y_source, preint_flow), flow_field
        else:
            return y_source, flow_field
//  [python]
def decode(self, z, interpolate=True):
        x = self.decoder(z)
        if interpolate:
            x = F.interpolate(x, scale_factor=self.interpolation_factor, mode='bilinear', align_corners=False)
        return x
//  [python]
def extract_features(self, images: torch.Tensor) -> torch.Tensor:
        """Extracting image features.

        Args:
            images (torch.Tensor): Images tensor.
        Returns:
            torch.Tensor: Vgg16 features of input images.
        """
        # image must passed in 'bgr'
        images = images[:, [2, 1, 0], ...]
        if self.use_tero_scirpt:
            images = images.to(torch.uint8)
            feature = self.vgg16(images, return_features=True)
        else:
            images = (images - 127.5) / 127.5
            batch = t.interpolate(images, size=(224, 224))
            before_fc = self.vgg16.features(batch)
            before_fc = before_fc.view(-1, 7 * 7 * 512)
            feature = self.vgg16.classifier[:4](before_fc)

        return feature
//  [python]
def extract_features(self, images: torch.Tensor) -> torch.Tensor:
        """Extracting image features.

        Args:
            images (torch.Tensor): Images tensor.
        Returns:
            torch.Tensor: Vgg16 features of input images.
        """
        # image must passed in 'bgr'
        images = images[:, [2, 1, 0], ...]
        if self.use_tero_scirpt:
            images = images.to(torch.uint8)
            feature = self.vgg16(images, return_features=True)
        else:
            images = (images - 127.5) / 127.5
            batch = F.interpolate(images, size=(224, 224))
            before_fc = self.vgg16.features(batch)
            before_fc = before_fc.view(-1, 7 * 7 * 512)
            feature = self.vgg16.classifier[:4](before_fc)

        return feature
//  [python]
def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] = laterals[i - 1] + t.interpolate(
                laterals[i], size=prev_shape, mode='nearest')

        # build outputs
        # part 1: from original levels
        inter_outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]

        # part 2: add bottom-up path
        for i in range(0, used_backbone_levels - 1):
            inter_outs[i + 1] = inter_outs[i + 1] + \
                                self.downsample_convs[i](inter_outs[i])

        outs = []
        outs.append(inter_outs[0])
        outs.extend([
            self.pafpn_convs[i - 1](inter_outs[i])
            for i in range(1, used_backbone_levels)
        ])

        # part 3: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    orig = inputs[self.backbone_end_level - 1]
                    outs.append(self.fpn_convs[used_backbone_levels](orig))
                elif self.add_extra_convs == 'on_lateral':
                    outs.append(self.fpn_convs[used_backbone_levels](
                        laterals[-1]))
                elif self.add_extra_convs == 'on_output':
                    outs.append(self.fpn_convs[used_backbone_levels](outs[-1]))
                else:
                    raise NotImplementedError
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)
//  [python]
def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] = laterals[i - 1] + F.interpolate(
                laterals[i], size=prev_shape, mode='nearest')

        # build outputs
        # part 1: from original levels
        inter_outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]

        # part 2: add bottom-up path
        for i in range(0, used_backbone_levels - 1):
            inter_outs[i + 1] = inter_outs[i + 1] + \
                                self.downsample_convs[i](inter_outs[i])

        outs = []
        outs.append(inter_outs[0])
        outs.extend([
            self.pafpn_convs[i - 1](inter_outs[i])
            for i in range(1, used_backbone_levels)
        ])

        # part 3: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    orig = inputs[self.backbone_end_level - 1]
                    outs.append(self.fpn_convs[used_backbone_levels](orig))
                elif self.add_extra_convs == 'on_lateral':
                    outs.append(self.fpn_convs[used_backbone_levels](
                        laterals[-1]))
                elif self.add_extra_convs == 'on_output':
                    outs.append(self.fpn_convs[used_backbone_levels](outs[-1]))
                else:
                    raise NotImplementedError
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)
//  [python]
def forward(self, x):
        B, C, H, W = x.shape
        if not self.fir:
            h = t.nn.functional.interpolate(x, (H * 2, W * 2), "nearest")
            if self.with_conv:
                h = self.Conv_0(h)
        else:
            if not self.with_conv:
                h = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
            else:
                h = self.Conv2d_0(x)

        return h
//  [python]
def forward(self, x):
        B, C, H, W = x.shape
        if not self.fir:
            h = F.interpolate(x, (H * 2, W * 2), "nearest")
            if self.with_conv:
                h = self.Conv_0(h)
        else:
            if not self.with_conv:
                h = up_or_down_sampling.upsample_2d(x, self.fir_kernel, factor=2)
            else:
                h = self.Conv2d_0(x)

        return h
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(self.conv_up1(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            if self.upscale == 4:
                x = self.lrelu(self.conv_up2(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            elif self.upscale == 8:
                x = self.lrelu(self.conv_up2(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
                x = self.lrelu(self.conv_up3(t.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, :H*self.upscale, :W*self.upscale]
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            if self.upscale == 4:
                x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            elif self.upscale == 8:
                x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
                x = self.lrelu(self.conv_up3(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, :H*self.upscale, :W*self.upscale]
//  [python]
def _upsample_mels(self, cond, audio_size, force_interpolate=False):
        if hasattr(self, 'upsample_net') and self.upsample_net is not None:
            cond = self.upsample_net(cond)
        if force_interpolate or (self.interpolation_required and not (cond.shape[2] == audio_size[2])):
            cond = t.interpolate(cond, size=audio_size[2], mode=self.upsample_mode, align_corners=True if self.upsample_mode == 'linear' else None)
        else:
            cond_len = cond.shape[2]
            audio_len = audio_size[2]
            pad_l = (cond_len-audio_len)//2
            pad_r = -(audio_len-cond_len)//2
            cond = cond[:, :, pad_l:-pad_r]
        return cond
//  [python]
def _upsample_mels(self, cond, audio_size, force_interpolate=False):
        if hasattr(self, 'upsample_net') and self.upsample_net is not None:
            cond = self.upsample_net(cond)
        if force_interpolate or (self.interpolation_required and not (cond.shape[2] == audio_size[2])):
            cond = F.interpolate(cond, size=audio_size[2], mode=self.upsample_mode, align_corners=True if self.upsample_mode == 'linear' else None)
        else:
            cond_len = cond.shape[2]
            audio_len = audio_size[2]
            pad_l = (cond_len-audio_len)//2
            pad_r = -(audio_len-cond_len)//2
            cond = cond[:, :, pad_l:-pad_r]
        return cond
//  [python]
def interpolate_pos_encoding(
        self, x,
    ):
        output_seq_len = x.shape[1]
        assert isPerfectSquare(output_seq_len - self.class_token_length)

        num_tok_output = output_seq_len - self.class_token_length
        num_tok_input = self.num_patches

        if num_tok_input == num_tok_output:
            return self.position_embeddings

        embed_tok = self.position_embeddings[: self.class_token_length]
        embed_grid = self.position_embeddings[self.class_token_length :]

        gs_new = int(math.sqrt(num_tok_output))
        gs_input = (self.num_patches_per_dim_h, self.num_patches_per_dim_w)

        embed_grid = embed_grid.transpose(0, 1).contiguous()
        embed_grid = embed_grid.reshape((1, -1, gs_input[0], gs_input[1]))
        embed_grid = embed_grid.float()
        scale_factor = (gs_new / gs_input[0], gs_new / gs_input[1])

        embed_grid = t.interpolate(embed_grid, scale_factor=scale_factor, mode="bicubic")

        embed_grid = embed_grid.reshape((-1, num_tok_output))
        embed_grid = embed_grid.transpose(0, 1).contiguous()

        return torch.cat((embed_tok, embed_grid), dim=0)
//  [python]
def interpolate_pos_encoding(
        self, x,
    ):
        output_seq_len = x.shape[1]
        assert isPerfectSquare(output_seq_len - self.class_token_length)

        num_tok_output = output_seq_len - self.class_token_length
        num_tok_input = self.num_patches

        if num_tok_input == num_tok_output:
            return self.position_embeddings

        embed_tok = self.position_embeddings[: self.class_token_length]
        embed_grid = self.position_embeddings[self.class_token_length :]

        gs_new = int(math.sqrt(num_tok_output))
        gs_input = (self.num_patches_per_dim_h, self.num_patches_per_dim_w)

        embed_grid = embed_grid.transpose(0, 1).contiguous()
        embed_grid = embed_grid.reshape((1, -1, gs_input[0], gs_input[1]))
        embed_grid = embed_grid.float()
        scale_factor = (gs_new / gs_input[0], gs_new / gs_input[1])

        embed_grid = F.interpolate(embed_grid, scale_factor=scale_factor, mode="bicubic")

        embed_grid = embed_grid.reshape((-1, num_tok_output))
        embed_grid = embed_grid.transpose(0, 1).contiguous()

        return torch.cat((embed_tok, embed_grid), dim=0)
//  [python]
def cal_appear_identity_sim(self, qobjs, kobjs):
        f0 = qobjs.get_feature()
        f1 = kobjs.get_feature()
        f1 = f1.to(f0)  # might be in cpu
        mask0 = qobjs.get_mask()
        mask1 = kobjs.get_mask()
        mask1 = mask1.to(mask0)  # might be in cpu
        mask0 = t.nn.functional.interpolate(mask0.unsqueeze(1), (f0.shape[2], f0.shape[3]), mode='bilinear',
                              align_corners=False).squeeze(1)
        mask1 = t.nn.functional.interpolate(mask1.unsqueeze(1), (f1.shape[2], f1.shape[3]), mode='bilinear',
                              align_corners=False).squeeze(1)
        sim = (f0 * f1 * mask0.unsqueeze(1) * mask1.unsqueeze(1)).sum([1, 2, 3]) / ((mask0 * mask1).sum([1, 2]) + 1e-6)
        return sim
//  [python]
def cal_appear_identity_sim(self, qobjs, kobjs):
        f0 = qobjs.get_feature()
        f1 = kobjs.get_feature()
        f1 = f1.to(f0)  # might be in cpu
        mask0 = qobjs.get_mask()
        mask1 = kobjs.get_mask()
        mask1 = mask1.to(mask0)  # might be in cpu
        mask0 = F.interpolate(mask0.unsqueeze(1), (f0.shape[2], f0.shape[3]), mode='bilinear',
                              align_corners=False).squeeze(1)
        mask1 = F.interpolate(mask1.unsqueeze(1), (f1.shape[2], f1.shape[3]), mode='bilinear',
                              align_corners=False).squeeze(1)
        sim = (f0 * f1 * mask0.unsqueeze(1) * mask1.unsqueeze(1)).sum([1, 2, 3]) / ((mask0 * mask1).sum([1, 2]) + 1e-6)
        return sim
//  [python]
def superres_T(self, T):
        # bilinear interpolation
        T_bchw = T.reshape(-1, self.corr_feat_height * self.corr_feat_width, self.corr_feat_height, self.corr_feat_width)
        T_bchw = t.nn.functional.interpolate(T_bchw, (self.corr_mask_height, self.corr_mask_width),
                               mode='bilinear', align_corners=False)
        T_b1hwc = T_bchw.reshape(-1, 1, self.corr_feat_height, self.corr_feat_width,
                                 self.corr_mask_height * self.corr_mask_width)
        T_b1hwc = t.nn.functional.interpolate(T_b1hwc,
                                (self.corr_mask_height, self.corr_mask_width, self.corr_mask_height * self.corr_mask_width),
                                mode='trilinear', align_corners=False)
        T_superres = T_b1hwc.reshape(-1, self.corr_mask_height * self.corr_mask_width,
                                     self.corr_mask_height * self.corr_mask_width) * \
                     (1.0 * self.corr_feat_height * self.corr_feat_width / self.corr_mask_height / self.corr_mask_width)

        return T_superres
//  [python]
def superres_T(self, T):
        # bilinear interpolation
        T_bchw = T.reshape(-1, self.corr_feat_height * self.corr_feat_width, self.corr_feat_height, self.corr_feat_width)
        T_bchw = F.interpolate(T_bchw, (self.corr_mask_height, self.corr_mask_width),
                               mode='bilinear', align_corners=False)
        T_b1hwc = T_bchw.reshape(-1, 1, self.corr_feat_height, self.corr_feat_width,
                                 self.corr_mask_height * self.corr_mask_width)
        T_b1hwc = F.interpolate(T_b1hwc,
                                (self.corr_mask_height, self.corr_mask_width, self.corr_mask_height * self.corr_mask_width),
                                mode='trilinear', align_corners=False)
        T_superres = T_b1hwc.reshape(-1, self.corr_mask_height * self.corr_mask_width,
                                     self.corr_mask_height * self.corr_mask_width) * \
                     (1.0 * self.corr_feat_height * self.corr_feat_width / self.corr_mask_height / self.corr_mask_width)

        return T_superres
//  [python]
def forward_single(self, x, idx, eval=False, upsampled_size=None):
        ins_kernel_feat = x
        # ins branch
        # concat coord
        x_range = torch.linspace(-1, 1, ins_kernel_feat.shape[-1], device=ins_kernel_feat.device)
        y_range = torch.linspace(-1, 1, ins_kernel_feat.shape[-2], device=ins_kernel_feat.device)
        y, x = torch.meshgrid(y_range, x_range)
        y = y.expand([ins_kernel_feat.shape[0], 1, -1, -1])
        x = x.expand([ins_kernel_feat.shape[0], 1, -1, -1])
        coord_feat = torch.cat([x, y], 1)
        ins_kernel_feat = torch.cat([ins_kernel_feat, coord_feat], 1)
        
        # kernel branch
        kernel_feat = ins_kernel_feat
        seg_num_grid = self.seg_num_grids[idx]
        kernel_feat = t.interpolate(kernel_feat, size=seg_num_grid, mode='bilinear')

        cate_feat = kernel_feat[:, :-2, :, :]

        kernel_feat = kernel_feat.contiguous()
        for i, kernel_layer in enumerate(self.kernel_convs):
            kernel_feat = kernel_layer(kernel_feat)
        kernel_pred = self.solo_kernel(kernel_feat)

        # cate branch
        cate_feat = cate_feat.contiguous()
        for i, cate_layer in enumerate(self.cate_convs):
            cate_feat = cate_layer(cate_feat)
        cate_pred = self.solo_cate(cate_feat)

        if eval:
            cate_pred = points_nms(cate_pred.sigmoid(), kernel=2).permute(0, 2, 3, 1)
        return cate_pred, kernel_pred
//  [python]
def forward_single(self, x, idx, eval=False, upsampled_size=None):
        ins_kernel_feat = x
        # ins branch
        # concat coord
        x_range = torch.linspace(-1, 1, ins_kernel_feat.shape[-1], device=ins_kernel_feat.device)
        y_range = torch.linspace(-1, 1, ins_kernel_feat.shape[-2], device=ins_kernel_feat.device)
        y, x = torch.meshgrid(y_range, x_range)
        y = y.expand([ins_kernel_feat.shape[0], 1, -1, -1])
        x = x.expand([ins_kernel_feat.shape[0], 1, -1, -1])
        coord_feat = torch.cat([x, y], 1)
        ins_kernel_feat = torch.cat([ins_kernel_feat, coord_feat], 1)
        
        # kernel branch
        kernel_feat = ins_kernel_feat
        seg_num_grid = self.seg_num_grids[idx]
        kernel_feat = F.interpolate(kernel_feat, size=seg_num_grid, mode='bilinear')

        cate_feat = kernel_feat[:, :-2, :, :]

        kernel_feat = kernel_feat.contiguous()
        for i, kernel_layer in enumerate(self.kernel_convs):
            kernel_feat = kernel_layer(kernel_feat)
        kernel_pred = self.solo_kernel(kernel_feat)

        # cate branch
        cate_feat = cate_feat.contiguous()
        for i, cate_layer in enumerate(self.cate_convs):
            cate_feat = cate_layer(cate_feat)
        cate_pred = self.solo_cate(cate_feat)

        if eval:
            cate_pred = points_nms(cate_pred.sigmoid(), kernel=2).permute(0, 2, 3, 1)
        return cate_pred, kernel_pred
//  [python]
def vis_corr(self, img_a_tensor, img_b_tensor, T, a_mask, b_mask, **img_norm_cfg):

        img_a = tensor2imgs(img_a_tensor, **img_norm_cfg)[0]
        img_b = tensor2imgs(img_b_tensor, **img_norm_cfg)[0]


        img_ab = np.zeros((img_a.shape[0], img_a.shape[1] + img_b.shape[1], 3), np.uint8)
        img_ab[:, :img_a.shape[1]] = img_a
        img_ab[:, img_a.shape[1]:] = img_b

        img_size = img_a.shape[0]

        assignment = T.argmax(1)
        size = int(assignment.shape[0] ** 0.5)
        assignment = assignment.reshape(size, size)

        a_mask = t.nn.functional.interpolate(a_mask.unsqueeze(0), (size, size), mode='bilinear', align_corners=False).squeeze()
        b_mask = t.nn.functional.interpolate(b_mask.unsqueeze(0), (size, size), mode='bilinear', align_corners=False).squeeze()

        scale = img_size / size

        for i in range(0,size,2):
            for j in range(0,size,2):
                x = assignment[i, j] % size
                y = assignment[i, j] // size
                if a_mask[i,j] > 0.5 and b_mask[x,y] > 0.5:
                    cv2.line(img_ab, (int(i * scale), int(j * scale)), (int((x + size) * scale), int(y * scale)),
                             color=tuple(self.color_panel[i*size + j].tolist()), thickness=5)

        self.vis_cnt += 1

        cv2.imwrite('corr_vis/{}.jpg'.format(self.vis_cnt), img_ab)
//  [python]
def vis_corr(self, img_a_tensor, img_b_tensor, T, a_mask, b_mask, **img_norm_cfg):

        img_a = tensor2imgs(img_a_tensor, **img_norm_cfg)[0]
        img_b = tensor2imgs(img_b_tensor, **img_norm_cfg)[0]


        img_ab = np.zeros((img_a.shape[0], img_a.shape[1] + img_b.shape[1], 3), np.uint8)
        img_ab[:, :img_a.shape[1]] = img_a
        img_ab[:, img_a.shape[1]:] = img_b

        img_size = img_a.shape[0]

        assignment = T.argmax(1)
        size = int(assignment.shape[0] ** 0.5)
        assignment = assignment.reshape(size, size)

        a_mask = F.interpolate(a_mask.unsqueeze(0), (size, size), mode='bilinear', align_corners=False).squeeze()
        b_mask = F.interpolate(b_mask.unsqueeze(0), (size, size), mode='bilinear', align_corners=False).squeeze()

        scale = img_size / size

        for i in range(0,size,2):
            for j in range(0,size,2):
                x = assignment[i, j] % size
                y = assignment[i, j] // size
                if a_mask[i,j] > 0.5 and b_mask[x,y] > 0.5:
                    cv2.line(img_ab, (int(i * scale), int(j * scale)), (int((x + size) * scale), int(y * scale)),
                             color=tuple(self.color_panel[i*size + j].tolist()), thickness=5)

        self.vis_cnt += 1

        cv2.imwrite('corr_vis/{}.jpg'.format(self.vis_cnt), img_ab)
//  [python]
def forward(self, ref, supp):
        """Forward function of SPyNet.

        This function computes the optical flow from ref to supp.

        Args:
            ref (Tensor): Reference image with shape of (n, 3, h, w).
            supp (Tensor): Supporting image with shape of (n, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (n, 2, h, w).
        """

        # upsize to a multiple of 32
        h, w = ref.shape[2:4]
        w_up = w if (w % 32) == 0 else 32 * (w // 32 + 1)
        h_up = h if (h % 32) == 0 else 32 * (h // 32 + 1)
        ref = t.interpolate(
            input=ref, size=(h_up, w_up), mode='bilinear', align_corners=False)
        supp = t.interpolate(
            input=supp,
            size=(h_up, w_up),
            mode='bilinear',
            align_corners=False)

        # compute flow, and resize back to the original resolution
        flow = t.interpolate(
            input=self.compute_flow(ref, supp),
            size=(h, w),
            mode='bilinear',
            align_corners=False)

        # adjust the flow values
        flow[:, 0, :, :] *= float(w) / float(w_up)
        flow[:, 1, :, :] *= float(h) / float(h_up)

        return flow
//  [python]
def forward(self, ref, supp):
        """Forward function of SPyNet.

        This function computes the optical flow from ref to supp.

        Args:
            ref (Tensor): Reference image with shape of (n, 3, h, w).
            supp (Tensor): Supporting image with shape of (n, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (n, 2, h, w).
        """

        # upsize to a multiple of 32
        h, w = ref.shape[2:4]
        w_up = w if (w % 32) == 0 else 32 * (w // 32 + 1)
        h_up = h if (h % 32) == 0 else 32 * (h // 32 + 1)
        ref = F.interpolate(
            input=ref, size=(h_up, w_up), mode='bilinear', align_corners=False)
        supp = F.interpolate(
            input=supp,
            size=(h_up, w_up),
            mode='bilinear',
            align_corners=False)

        # compute flow, and resize back to the original resolution
        flow = F.interpolate(
            input=self.compute_flow(ref, supp),
            size=(h, w),
            mode='bilinear',
            align_corners=False)

        # adjust the flow values
        flow[:, 0, :, :] *= float(w) / float(w_up)
        flow[:, 1, :, :] *= float(h) / float(h_up)

        return flow
//  [python]
def compute_flow(self, ref, supp):
        """Compute flow from ref to supp.

        Note that in this function, the images are already resized to a
        multiple of 32.

        Args:
            ref (Tensor): Reference image with shape of (n, 3, h, w).
            supp (Tensor): Supporting image with shape of (n, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (n, 2, h, w).
        """
        n, _, h, w = ref.size()

        # normalize the input images
        ref = [(ref - self.mean) / self.std]
        supp = [(supp - self.mean) / self.std]

        # generate downsampled frames
        for level in range(5):
            ref.append(
                F.avg_pool2d(
                    input=ref[-1],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
            supp.append(
                F.avg_pool2d(
                    input=supp[-1],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
        ref = ref[::-1]
        supp = supp[::-1]

        # flow computation
        flow = ref[0].new_zeros(n, 2, h // 32, w // 32)
        for level in range(len(ref)):
            if level == 0:
                flow_up = flow
            else:
                flow_up = t.upsample(
                    input=flow,
                    scale_factor=2,
                    mode='bilinear') * 2.0

            # add the residue to the upsampled flow
            flow = flow_up + self.basic_module[level](
                torch.cat([
                    ref[level],
                    flow_warp(
                        supp[level],
                        flow_up.permute(0, 2, 3, 1),
                        padding_mode='border'), flow_up
                ], 1))

        return flow
//  [python]
def compute_flow(self, ref, supp):
        """Compute flow from ref to supp.

        Note that in this function, the images are already resized to a
        multiple of 32.

        Args:
            ref (Tensor): Reference image with shape of (n, 3, h, w).
            supp (Tensor): Supporting image with shape of (n, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (n, 2, h, w).
        """
        n, _, h, w = ref.size()

        # normalize the input images
        ref = [(ref - self.mean) / self.std]
        supp = [(supp - self.mean) / self.std]

        # generate downsampled frames
        for level in range(5):
            ref.append(
                F.avg_pool2d(
                    input=ref[-1],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
            supp.append(
                F.avg_pool2d(
                    input=supp[-1],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
        ref = ref[::-1]
        supp = supp[::-1]

        # flow computation
        flow = ref[0].new_zeros(n, 2, h // 32, w // 32)
        for level in range(len(ref)):
            if level == 0:
                flow_up = flow
            else:
                flow_up = F.interpolate(
                    input=flow,
                    scale_factor=2,
                    mode='bilinear',
                    align_corners=True) * 2.0

            # add the residue to the upsampled flow
            flow = flow_up + self.basic_module[level](
                torch.cat([
                    ref[level],
                    flow_warp(
                        supp[level],
                        flow_up.permute(0, 2, 3, 1),
                        padding_mode='border'), flow_up
                ], 1))

        return flow
//  [python]
def slice(self, tensor, slice_shape):
        t_shape = tensor.shape
        img_h, img_w = slice_shape
        if img_w!=t_shape[-1] or img_h!=t_shape[2]:
            new_tensor = t.imresize(tensor,(img_h, img_w),interpolation='bicubic')

        else:
            new_tensor=tensor
        # tensor[..., :height, :width]
        return new_tensor
//  [python]
def slice(self, tensor, slice_shape):
        t_shape = tensor.shape
        img_h, img_w = slice_shape
        if img_w!=t_shape[-1] or img_h!=t_shape[2]:
            new_tensor = F.interpolate(
                tensor, size=(img_h, img_w), mode='bicubic',align_corners=False)

        else:
            new_tensor=tensor
        # tensor[..., :height, :width]
        return new_tensor
//  [python]
def resize_input(self,tensor):
        t_shape = tensor.shape
        if t_shape[2] % 8 != 0 or t_shape[3] % 8 != 0:
            img_w= ((t_shape[3]// 8) + 1) * 8
            img_h = ((t_shape[2] // 8) + 1) * 8
            new_tensor = t.nn.functional.interpolate(
                tensor, size=(img_h, img_w), mode='bicubic', align_corners=False)
        else:
            new_tensor = tensor
        return new_tensor
//  [python]
def resize_input(self,tensor):
        t_shape = tensor.shape
        if t_shape[2] % 8 != 0 or t_shape[3] % 8 != 0:
            img_w= ((t_shape[3]// 8) + 1) * 8
            img_h = ((t_shape[2] // 8) + 1) * 8
            new_tensor = F.interpolate(
                tensor, size=(img_h, img_w), mode='bicubic', align_corners=False)
        else:
            new_tensor = tensor
        return new_tensor
//  [python]
def forward(self, features):
        res2 = features['res2']
        res3 = features['res3']
        res4 = features['res4']
        res5 = features['res5']

        pred = OrderedDict()

        res5_feature = self.res5_fuse(res5)
        for key in self.class_key:
            pred[key + '-5'] = self.classifiers[key + '-5'](res5_feature)

        res4_feature = self.res4_fuse(res4, res5_feature)
        #res4_feature = self.res4_ocr(res4_feature, pred['semantic-5'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-4'] = self.classifiers[key + '-4'](res4_feature)
            else:
                offset_pred = t.interpolate(pred['offset-5'],
                                            size=res4_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-4'](torch.cat([res4_feature, offset_pred], dim=1))
                pred[key + '-4'] = offset_pred + residual

        res3_feature = self.res3_fuse(res3, res4_feature)
        #res3_feature = self.res3_ocr(res3_feature, pred['semantic-4'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-3'] = self.classifiers[key + '-3'](res3_feature)
            else:
                offset_pred = t.interpolate(pred['offset-4'],
                                            size=res3_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-3'](torch.cat([res3_feature, offset_pred], dim=1))
                pred[key + '-3'] = offset_pred + residual

        res2_feature = self.res2_fuse(res2, res3_feature)
        res2_feature = self.res2_ocr(res2_feature, pred['semantic-3'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-2'] = self.classifiers[key + '-2'](res2_feature)
            else:
                offset_pred = t.interpolate(pred['offset-3'],
                                            size=res2_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-2'](torch.cat([res2_feature, offset_pred], dim=1))
                pred[key + '-2'] = offset_pred + residual

        return pred
//  [python]
def forward(self, features):
        res2 = features['res2']
        res3 = features['res3']
        res4 = features['res4']
        res5 = features['res5']

        pred = OrderedDict()

        res5_feature = self.res5_fuse(res5)
        for key in self.class_key:
            pred[key + '-5'] = self.classifiers[key + '-5'](res5_feature)

        res4_feature = self.res4_fuse(res4, res5_feature)
        #res4_feature = self.res4_ocr(res4_feature, pred['semantic-5'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-4'] = self.classifiers[key + '-4'](res4_feature)
            else:
                offset_pred = F.interpolate(pred['offset-5'],
                                            size=res4_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-4'](torch.cat([res4_feature, offset_pred], dim=1))
                pred[key + '-4'] = offset_pred + residual

        res3_feature = self.res3_fuse(res3, res4_feature)
        #res3_feature = self.res3_ocr(res3_feature, pred['semantic-4'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-3'] = self.classifiers[key + '-3'](res3_feature)
            else:
                offset_pred = F.interpolate(pred['offset-4'],
                                            size=res3_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-3'](torch.cat([res3_feature, offset_pred], dim=1))
                pred[key + '-3'] = offset_pred + residual

        res2_feature = self.res2_fuse(res2, res3_feature)
        res2_feature = self.res2_ocr(res2_feature, pred['semantic-3'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-2'] = self.classifiers[key + '-2'](res2_feature)
            else:
                offset_pred = F.interpolate(pred['offset-3'],
                                            size=res2_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-2'](torch.cat([res2_feature, offset_pred], dim=1))
                pred[key + '-2'] = offset_pred + residual

        return pred
//  [python]
def get_mask(self,x_t, maps, alpha, use_pool):
        k = 1
        maps = (maps * alpha).sum(-1).mean(1)
        if use_pool:
            maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
        mask = t.nn.functional.interpolate(maps, size=(x_t.shape[2:]))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.th[1-int(use_pool)])
        mask = mask[:1] + mask
        return mask
//  [python]
def get_mask(self,x_t, maps, alpha, use_pool):
        k = 1
        maps = (maps * alpha).sum(-1).mean(1)
        if use_pool:
            maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
        mask = nnf.interpolate(maps, size=(x_t.shape[2:]))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.th[1-int(use_pool)])
        mask = mask[:1] + mask
        return mask
//  [python]
def forward(self, features):
        coord_features = self.compute_coordinates(features)
        features = torch.cat([coord_features, features], dim=1)
        pred_logits, pred_kernel, pred_scores, iam = self.inst_branch(features)
        mask_features = self.mask_branch(features)

        N = pred_kernel.shape[1]
        # mask_features: BxCxHxW
        B, C, H, W = mask_features.shape
        pred_masks = torch.bmm(pred_kernel,
                               mask_features.view(B, C,
                                                  H * W)).view(B, N, H, W)

        pred_masks = F.upsample(
            pred_masks,
            scale_factor=self.scale_factor,
            mode='bilinear')

        output = {
            'pred_logits': pred_logits,
            'pred_masks': pred_masks,
            'pred_scores': pred_scores,
        }

        if self.output_iam:
            iam = F.upsample(
                iam,
                scale_factor=self.scale_factor,
                mode='bilinear')
            output['pred_iam'] = iam

        return output
//  [python]
def forward(self, features):
        coord_features = self.compute_coordinates(features)
        features = torch.cat([coord_features, features], dim=1)
        pred_logits, pred_kernel, pred_scores, iam = self.inst_branch(features)
        mask_features = self.mask_branch(features)

        N = pred_kernel.shape[1]
        # mask_features: BxCxHxW
        B, C, H, W = mask_features.shape
        pred_masks = torch.bmm(pred_kernel,
                               mask_features.view(B, C,
                                                  H * W)).view(B, N, H, W)

        pred_masks = F.interpolate(
            pred_masks,
            scale_factor=self.scale_factor,
            mode='bilinear',
            align_corners=False)

        output = {
            'pred_logits': pred_logits,
            'pred_masks': pred_masks,
            'pred_scores': pred_scores,
        }

        if self.output_iam:
            iam = F.interpolate(
                iam,
                scale_factor=self.scale_factor,
                mode='bilinear',
                align_corners=False)
            output['pred_iam'] = iam

        return output
//  [python]
def resize_feats(self, x: Tuple[Tensor]) -> List[Tensor]:
        """Downsample the first feat and upsample last feat in feats.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            list[Tensor]: Features after resizing, each is a 4D-tensor.
        """
        out = []
        for i in range(len(x)):
            if i == 0:
                out.append(
                    t.interpolate(x[0], scale_factor=0.5, mode='bilinear'))
            elif i == len(x) - 1:
                out.append(
                    t.interpolate(
                        x[i], size=x[i - 1].shape[-2:], mode='bilinear'))
            else:
                out.append(x[i])
        return out
//  [python]
def resize_feats(self, x: Tuple[Tensor]) -> List[Tensor]:
        """Downsample the first feat and upsample last feat in feats.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            list[Tensor]: Features after resizing, each is a 4D-tensor.
        """
        out = []
        for i in range(len(x)):
            if i == 0:
                out.append(
                    F.interpolate(x[0], scale_factor=0.5, mode='bilinear'))
            elif i == len(x) - 1:
                out.append(
                    F.interpolate(
                        x[i], size=x[i - 1].shape[-2:], mode='bilinear'))
            else:
                out.append(x[i])
        return out
//  [python]
def forward(
        self,
        x,
        mask,
        feats,
        img,
        ws,
        fname=None,
        force_fp32=False,
        fused_modconv=None,
        **layer_kwargs,
    ):
        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
        dtype = torch.float32
        memory_format = (
            torch.channels_last
            if self.channels_last and not force_fp32
            else torch.contiguous_format
        )
        if fused_modconv is None:
            fused_modconv = (not self.training) and (
                dtype == torch.float32 or int(x.shape[0]) == 1
            )

        x = x.to(dtype=dtype, memory_format=memory_format)
        x_skip = (
            feats[self.resolution].clone().to(dtype=dtype, memory_format=memory_format)
        )

        # Main layers.
        if self.in_channels == 0:
            x = self.conv1(x, ws[1], fused_modconv=fused_modconv, **layer_kwargs)
        elif self.architecture == "resnet":
            y = self.skip(x, gain=np.sqrt(0.5))
            x = self.conv0(
                x, ws[0].clone(), fused_modconv=fused_modconv, **layer_kwargs
            )
            if len(self.ffc_skip) > 0:
                mask = F.upsample(
                    mask,
                    size=x_skip.shape[2:],
                    mode='nearest'
                )
                z = x + x_skip
                for fres in self.ffc_skip:
                    z = fres(z, mask)
                x = x + z
            else:
                x = x + x_skip
            x = self.conv1(
                x,
                ws[1].clone(),
                fused_modconv=fused_modconv,
                gain=np.sqrt(0.5),
                **layer_kwargs,
            )
            x = y.add_(x)
        else:
            x = self.conv0(
                x, ws[0].clone(), fused_modconv=fused_modconv, **layer_kwargs
            )
            if len(self.ffc_skip) > 0:
                mask = F.upsample(
                    mask,
                    size=x_skip.shape[2:],
                    mode='nearest'
                )
                z = x + x_skip
                for fres in self.ffc_skip:
                    z = fres(z, mask)
                x = x + z
            else:
                x = x + x_skip
            x = self.conv1(
                x, ws[1].clone(), fused_modconv=fused_modconv, **layer_kwargs
            )
        # ToRGB.
        if img is not None:
            img = upsample2d(img, self.resample_filter)
        if self.is_last or self.architecture == "skip":
            y = self.torgb(x, ws[2].clone(), fused_modconv=fused_modconv)
            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
            img = img.add_(y) if img is not None else y

        x = x.to(dtype=dtype)
        assert x.dtype == dtype
        assert img is None or img.dtype == torch.float32
        return x, img