jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _mask_point_forward_test(self, x, rois, label_pred, mask_pred,
                                 img_metas):
        """Mask refining process with point head in testing."""
        refined_mask_pred = mask_pred.clone()
        for subdivision_step in range(self.test_cfg.subdivision_steps):
            refined_mask_pred = t.nn.functional.upsample(
                refined_mask_pred,
                scale_factor=self.test_cfg.scale_factor,
                mode='bilinear')
            # If `subdivision_num_points` is larger or equal to the
            # resolution of the next step, then we can skip this step
            num_rois, channels, mask_height, mask_width = \
                refined_mask_pred.shape
            if (self.test_cfg.subdivision_num_points >=
                    self.test_cfg.scale_factor**2 * mask_height * mask_width
                    and
                    subdivision_step < self.test_cfg.subdivision_steps - 1):
                continue
            point_indices, rel_roi_points = \
                self.point_head.get_roi_rel_points_test(
                    refined_mask_pred, label_pred, cfg=self.test_cfg)
            fine_grained_point_feats = self._get_fine_grained_point_feats(
                x, rois, rel_roi_points, img_metas)
            coarse_point_feats = point_sample(mask_pred, rel_roi_points)
            mask_point_pred = self.point_head(fine_grained_point_feats,
                                              coarse_point_feats)

            point_indices = point_indices.unsqueeze(1).expand(-1, channels, -1)
            refined_mask_pred = refined_mask_pred.reshape(
                num_rois, channels, mask_height * mask_width)
            refined_mask_pred = refined_mask_pred.scatter_(
                2, point_indices, mask_point_pred)
            refined_mask_pred = refined_mask_pred.view(num_rois, channels,
                                                       mask_height, mask_width)

        return refined_mask_pred
//  [python]
def _mask_point_forward_test(self, x, rois, label_pred, mask_pred,
                                 img_metas):
        """Mask refining process with point head in testing."""
        refined_mask_pred = mask_pred.clone()
        for subdivision_step in range(self.test_cfg.subdivision_steps):
            refined_mask_pred = F.interpolate(
                refined_mask_pred,
                scale_factor=self.test_cfg.scale_factor,
                mode='bilinear',
                align_corners=False)
            # If `subdivision_num_points` is larger or equal to the
            # resolution of the next step, then we can skip this step
            num_rois, channels, mask_height, mask_width = \
                refined_mask_pred.shape
            if (self.test_cfg.subdivision_num_points >=
                    self.test_cfg.scale_factor**2 * mask_height * mask_width
                    and
                    subdivision_step < self.test_cfg.subdivision_steps - 1):
                continue
            point_indices, rel_roi_points = \
                self.point_head.get_roi_rel_points_test(
                    refined_mask_pred, label_pred, cfg=self.test_cfg)
            fine_grained_point_feats = self._get_fine_grained_point_feats(
                x, rois, rel_roi_points, img_metas)
            coarse_point_feats = point_sample(mask_pred, rel_roi_points)
            mask_point_pred = self.point_head(fine_grained_point_feats,
                                              coarse_point_feats)

            point_indices = point_indices.unsqueeze(1).expand(-1, channels, -1)
            refined_mask_pred = refined_mask_pred.reshape(
                num_rois, channels, mask_height * mask_width)
            refined_mask_pred = refined_mask_pred.scatter_(
                2, point_indices, mask_point_pred)
            refined_mask_pred = refined_mask_pred.view(num_rois, channels,
                                                       mask_height, mask_width)

        return refined_mask_pred
//  [python]
def forward(self, theta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

        backcast = theta[:, :self.backcast_size]
        knots = theta[:, self.backcast_size:]

        if self.interpolation_mode in ['nearest','linear']:
            knots = knots[:,None,:]
            forecast = t.interpolate(knots, size=self.forecast_size, mode=self.interpolation_mode)
            forecast = forecast[:,0,:]
        elif 'cubic' in self.interpolation_mode:
            batch_size = len(backcast)
            knots = knots[:,None,None,:]
            forecast = torch.zeros((len(knots), self.forecast_size)).to(knots.device)
            n_batches = int(np.ceil(len(knots)/batch_size))
            for i in range(n_batches):
                forecast_i = t.interpolate(knots[i*batch_size:(i+1)*batch_size], size=self.forecast_size, mode='bicubic')
                forecast[i*batch_size:(i+1)*batch_size] += forecast_i[:,0,0,:]

        return backcast, forecast
//  [python]
def forward(self, theta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

        backcast = theta[:, :self.backcast_size]
        knots = theta[:, self.backcast_size:]

        if self.interpolation_mode in ['nearest','linear']:
            knots = knots[:,None,:]
            forecast = F.interpolate(knots, size=self.forecast_size, mode=self.interpolation_mode)
            forecast = forecast[:,0,:]
        elif 'cubic' in self.interpolation_mode:
            batch_size = len(backcast)
            knots = knots[:,None,None,:]
            forecast = torch.zeros((len(knots), self.forecast_size)).to(knots.device)
            n_batches = int(np.ceil(len(knots)/batch_size))
            for i in range(n_batches):
                forecast_i = F.interpolate(knots[i*batch_size:(i+1)*batch_size], size=self.forecast_size, mode='bicubic')
                forecast[i*batch_size:(i+1)*batch_size] += forecast_i[:,0,0,:]

        return backcast, forecast
//  [python]
def __call__(self, sample):
        width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])

        # resize sample
        sample["image"] = cv2.resize(
            sample["image"],
            (width, height),
            interpolation=self.__image_interpolation_method,
        )

        if self.__resize_target:
            if "disparity" in sample:
                sample["disparity"] = cv2.resize(
                    sample["disparity"],
                    (width, height),
                    interpolation=cv2.INTER_NEAREST,
                )

            if "depth" in sample:
                sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)

            if "semseg_mask" in sample:
                sample["semseg_mask"] = cv2.resize(
                    sample["semseg_mask"], (width, height), interpolation=cv2.INTER_NEAREST
                )
                # sample["semseg_mask"] = F.interpolate(
                #     torch.from_numpy(sample["semseg_mask"]).float()[None, None, ...], (height, width), mode="nearest"
                # ).numpy()[0, 0]

            if "mask" in sample:
                sample["mask"] = cv2.resize(
                    sample["mask"].astype(np.float32),
                    (width, height),
                    interpolation=cv2.INTER_NEAREST,
                )
                # sample["mask"] = sample["mask"].astype(bool)

        # print(sample['image'].shape, sample['depth'].shape)
        return sample
//  [python]
def __call__(self, sample):
        width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])

        # resize sample
        sample["image"] = cv2.resize(
            sample["image"],
            (width, height),
            interpolation=self.__image_interpolation_method,
        )

        if self.__resize_target:
            if "disparity" in sample:
                sample["disparity"] = cv2.resize(
                    sample["disparity"],
                    (width, height),
                    interpolation=cv2.INTER_NEAREST,
                )

            if "depth" in sample:
                sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)

            if "semseg_mask" in sample:
                # sample["semseg_mask"] = cv2.resize(
                #     sample["semseg_mask"], (width, height), interpolation=cv2.INTER_NEAREST
                # )
                sample["semseg_mask"] = F.interpolate(
                    torch.from_numpy(sample["semseg_mask"]).float()[None, None, ...], (height, width), mode="nearest"
                ).numpy()[0, 0]

            if "mask" in sample:
                sample["mask"] = cv2.resize(
                    sample["mask"].astype(np.float32),
                    (width, height),
                    interpolation=cv2.INTER_NEAREST,
                )
                # sample["mask"] = sample["mask"].astype(bool)

        # print(sample['image'].shape, sample['depth'].shape)
        return sample
//  [python]
def forward(self, x):
        H, W = x.size()[2:]
        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
        feat_fuse = self.ffm(feat_sp, feat_cp8)

        feat_out = self.conv_out(feat_fuse)
        feat_out16 = self.conv_out16(feat_cp8)
        feat_out32 = self.conv_out32(feat_cp16)

        feat_out = F.upsample(feat_out, size=(H, W), mode='bilinear')
        feat_out16 = F.upsample(feat_out16, size=(H, W), mode='bilinear')
        feat_out32 = F.upsample(feat_out32, size=(H, W), mode='bilinear')
        return feat_out, feat_out16, feat_out32
//  [python]
def forward(self, x):
        H, W = x.size()[2:]
        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
        feat_fuse = self.ffm(feat_sp, feat_cp8)

        feat_out = self.conv_out(feat_fuse)
        feat_out16 = self.conv_out16(feat_cp8)
        feat_out32 = self.conv_out32(feat_cp16)

        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
        return feat_out, feat_out16, feat_out32
//  [python]
def forward(self, x):
        H0, W0 = x.size()[2:]
        feat8, feat16, feat32 = self.resnet(x)
        H8, W8 = feat8.size()[2:]
        H16, W16 = feat16.size()[2:]
        H32, W32 = feat32.size()[2:]

        avg = F.avg_pool2d(feat32, feat32.size()[2:])
        avg = self.conv_avg(avg)
        avg_up = F.upsample(avg, (H32, W32), mode='nearest')

        feat32_arm = self.arm32(feat32)
        feat32_sum = feat32_arm + avg_up
        feat32_up = F.upsample(feat32_sum, (H16, W16), mode='nearest')
        feat32_up = self.conv_head32(feat32_up)

        feat16_arm = self.arm16(feat16)
        feat16_sum = feat16_arm + feat32_up
        feat16_up = F.upsample(feat16_sum, (H8, W8), mode='nearest')
        feat16_up = self.conv_head16(feat16_up)

        return feat8, feat16_up, feat32_up
//  [python]
def forward(self, x):
        H0, W0 = x.size()[2:]
        feat8, feat16, feat32 = self.resnet(x)
        H8, W8 = feat8.size()[2:]
        H16, W16 = feat16.size()[2:]
        H32, W32 = feat32.size()[2:]

        avg = F.avg_pool2d(feat32, feat32.size()[2:])
        avg = self.conv_avg(avg)
        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')

        feat32_arm = self.arm32(feat32)
        feat32_sum = feat32_arm + avg_up
        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
        feat32_up = self.conv_head32(feat32_up)

        feat16_arm = self.arm16(feat16)
        feat16_sum = feat16_arm + feat32_up
        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
        feat16_up = self.conv_head16(feat16_up)

        return feat8, feat16_up, feat32_up
//  [python]
def _make_image_pyramid(self, intensity_img: torch.Tensor, depth_img: torch.Tensor):
        d0_w, d0_h = intensity_img.size(1), intensity_img.size(0)
        d1_w, d1_h = d0_w // 2, d0_h // 2
        d2_w, d2_h = d1_w // 2, d1_h // 2
        d0_intensity = intensity_img.view(1, 1, d0_h, d0_w)
        d0_depth = depth_img.view(1, 1, d0_h, d0_w)
        d1_intensity = t(d0_intensity, (d1_h, d1_w), mode='bilinear')
        d1_depth = t(d0_depth, (d1_h, d1_w), mode='nearest')
        d2_intensity = t(d1_intensity, (d2_h, d2_w), mode='bilinear')
        d2_depth = t(d1_depth, (d2_h, d2_w), mode='nearest')
        d0_gradient = gradient_xy(d0_intensity.squeeze(0).squeeze(0))
        d1_gradient = gradient_xy(d1_intensity.squeeze(0).squeeze(0))
        d2_gradient = gradient_xy(d2_intensity.squeeze(0).squeeze(0))
        return [t.squeeze(0).squeeze(0) for t in [d0_intensity, d1_intensity, d2_intensity]], \
               [t.squeeze(0).squeeze(0) for t in [d0_depth, d1_depth, d2_depth]], \
               [d0_gradient, d1_gradient, d2_gradient]
//  [python]
def _make_image_pyramid(self, intensity_img: torch.Tensor, depth_img: torch.Tensor):
        d0_w, d0_h = intensity_img.size(1), intensity_img.size(0)
        d1_w, d1_h = d0_w // 2, d0_h // 2
        d2_w, d2_h = d1_w // 2, d1_h // 2
        d0_intensity = intensity_img.view(1, 1, d0_h, d0_w)
        d0_depth = depth_img.view(1, 1, d0_h, d0_w)
        d1_intensity = torch.nn.functional.interpolate(d0_intensity, (d1_h, d1_w), mode='bilinear')
        d1_depth = torch.nn.functional.interpolate(d0_depth, (d1_h, d1_w), mode='nearest')
        d2_intensity = torch.nn.functional.interpolate(d1_intensity, (d2_h, d2_w), mode='bilinear')
        d2_depth = torch.nn.functional.interpolate(d1_depth, (d2_h, d2_w), mode='nearest')
        d0_gradient = gradient_xy(d0_intensity.squeeze(0).squeeze(0))
        d1_gradient = gradient_xy(d1_intensity.squeeze(0).squeeze(0))
        d2_gradient = gradient_xy(d2_intensity.squeeze(0).squeeze(0))
        return [t.squeeze(0).squeeze(0) for t in [d0_intensity, d1_intensity, d2_intensity]], \
               [t.squeeze(0).squeeze(0) for t in [d0_depth, d1_depth, d2_depth]], \
               [d0_gradient, d1_gradient, d2_gradient]
//  [python]
def forward(self, inp):
        """Get Inception feature maps

        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)

        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        outp = []
        x = inp

        if self.resize_input:
            x = t.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)

        if self.normalize_input:
            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)

            if idx == self.last_needed_block:
                break

        return outp
//  [python]
def forward(self, inp):
        """Get Inception feature maps

        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)

        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        outp = []
        x = inp

        if self.resize_input:
            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)

        if self.normalize_input:
            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)

            if idx == self.last_needed_block:
                break

        return outp
//  [python]
def forward(self, x):
        half_check = False
        if x.type() == "torch.cuda.HalfTensor":
            # half only works on gpu anyway
            half_check = True

        batch = x.shape[0]

        if self.spatial_scale_factor is not None:
            orig_size = x.shape[-2:]
            x = t.interpolate(
                x,
                scale_factor=self.spatial_scale_factor,
                mode=self.spatial_scale_mode,
                align_corners=False,
            )

        # (batch, c, h, w/2+1, 2)
        fft_dim = (-3, -2, -1) if self.ffc3d else (-2, -1)
        if half_check == True:
            ffted = torch.fft.rfftn(
                x.float(), dim=fft_dim, norm=self.fft_norm
            )  # .type(torch.cuda.HalfTensor)
        else:
            ffted = torch.fft.rfftn(x, dim=fft_dim, norm=self.fft_norm)

        ffted = torch.stack((ffted.real, ffted.imag), dim=-1)
        ffted = ffted.permute(0, 1, 4, 2, 3).contiguous()  # (batch, c, 2, h, w/2+1)
        ffted = ffted.view(
            (
                batch,
                -1,
            )
            + ffted.size()[3:]
        )

        if self.spectral_pos_encoding:
            height, width = ffted.shape[-2:]
            coords_vert = (
                torch.linspace(0, 1, height)[None, None, :, None]
                .expand(batch, 1, height, width)
                .to(ffted)
            )
            coords_hor = (
                torch.linspace(0, 1, width)[None, None, None, :]
                .expand(batch, 1, height, width)
                .to(ffted)
            )
            ffted = torch.cat((coords_vert, coords_hor, ffted), dim=1)

        if self.use_se:
            ffted = self.se(ffted)

        if half_check == True:
            ffted = self.conv_layer(ffted.half())  # (batch, c*2, h, w/2+1)
        else:
            ffted = self.conv_layer(
                ffted
            )  # .type(torch.cuda.FloatTensor)  # (batch, c*2, h, w/2+1)

        ffted = self.relu(self.bn(ffted))
        # forcing to be always float
        ffted = ffted.float()

        ffted = (
            ffted.view(
                (
                    batch,
                    -1,
                    2,
                )
                + ffted.size()[2:]
            )
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )  # (batch,c, t, h, w/2+1, 2)

        ffted = torch.complex(ffted[..., 0], ffted[..., 1])

        ifft_shape_slice = x.shape[-3:] if self.ffc3d else x.shape[-2:]
        output = torch.fft.irfftn(
            ffted, s=ifft_shape_slice, dim=fft_dim, norm=self.fft_norm
        )

        if half_check == True:
            output = output.half()

        if self.spatial_scale_factor is not None:
            output = t.interpolate(
                output,
                size=orig_size,
                mode=self.spatial_scale_mode,
                align_corners=False,
            )

        return output
//  [python]
def forward(self, x):
        half_check = False
        if x.type() == "torch.cuda.HalfTensor":
            # half only works on gpu anyway
            half_check = True

        batch = x.shape[0]

        if self.spatial_scale_factor is not None:
            orig_size = x.shape[-2:]
            x = F.interpolate(
                x,
                scale_factor=self.spatial_scale_factor,
                mode=self.spatial_scale_mode,
                align_corners=False,
            )

        # (batch, c, h, w/2+1, 2)
        fft_dim = (-3, -2, -1) if self.ffc3d else (-2, -1)
        if half_check == True:
            ffted = torch.fft.rfftn(
                x.float(), dim=fft_dim, norm=self.fft_norm
            )  # .type(torch.cuda.HalfTensor)
        else:
            ffted = torch.fft.rfftn(x, dim=fft_dim, norm=self.fft_norm)

        ffted = torch.stack((ffted.real, ffted.imag), dim=-1)
        ffted = ffted.permute(0, 1, 4, 2, 3).contiguous()  # (batch, c, 2, h, w/2+1)
        ffted = ffted.view(
            (
                batch,
                -1,
            )
            + ffted.size()[3:]
        )

        if self.spectral_pos_encoding:
            height, width = ffted.shape[-2:]
            coords_vert = (
                torch.linspace(0, 1, height)[None, None, :, None]
                .expand(batch, 1, height, width)
                .to(ffted)
            )
            coords_hor = (
                torch.linspace(0, 1, width)[None, None, None, :]
                .expand(batch, 1, height, width)
                .to(ffted)
            )
            ffted = torch.cat((coords_vert, coords_hor, ffted), dim=1)

        if self.use_se:
            ffted = self.se(ffted)

        if half_check == True:
            ffted = self.conv_layer(ffted.half())  # (batch, c*2, h, w/2+1)
        else:
            ffted = self.conv_layer(
                ffted
            )  # .type(torch.cuda.FloatTensor)  # (batch, c*2, h, w/2+1)

        ffted = self.relu(self.bn(ffted))
        # forcing to be always float
        ffted = ffted.float()

        ffted = (
            ffted.view(
                (
                    batch,
                    -1,
                    2,
                )
                + ffted.size()[2:]
            )
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )  # (batch,c, t, h, w/2+1, 2)

        ffted = torch.complex(ffted[..., 0], ffted[..., 1])

        ifft_shape_slice = x.shape[-3:] if self.ffc3d else x.shape[-2:]
        output = torch.fft.irfftn(
            ffted, s=ifft_shape_slice, dim=fft_dim, norm=self.fft_norm
        )

        if half_check == True:
            output = output.half()

        if self.spatial_scale_factor is not None:
            output = F.interpolate(
                output,
                size=orig_size,
                mode=self.spatial_scale_mode,
                align_corners=False,
            )

        return output
//  [python]
def forward(self, x):
        e0 = self.en_block0(x)
        e1 = self.en_block1(e0)
        e2 = self.en_block2(e1)
        e3 = self.en_block3(e2)
        e4 = self.en_block4(e3)
        
        d4 = self.de_block4(e4)
        d4 = t.interpolate(d4, size=e3.size()[2:], mode='bilinear', align_corners=True)
        c4 = torch.cat((d4,e3),1)

        d3 = self.de_block3(c4)
        d3 = t.interpolate(d3, size=e2.size()[2:], mode='bilinear', align_corners=True)
        c3 = torch.cat((d3,e2),1)

        d2 = self.de_block2(c3)
        d2 = t.interpolate(d2, size=e1.size()[2:], mode='bilinear', align_corners=True)
        c2 =torch.cat((d2,e1),1)

        d1 = self.de_block1(c2)
        d1 = t.interpolate(d1, size=e0.size()[2:], mode='bilinear', align_corners=True)
        c1 = torch.cat((d1,e0),1)
        y = self.de_block0(c1)
        
        return y
//  [python]
def forward(self, x):
        e0 = self.en_block0(x)
        e1 = self.en_block1(e0)
        e2 = self.en_block2(e1)
        e3 = self.en_block3(e2)
        e4 = self.en_block4(e3)
        
        d4 = self.de_block4(e4)
        d4 = F.interpolate(d4, size=e3.size()[2:], mode='bilinear', align_corners=True)
        c4 = torch.cat((d4,e3),1)

        d3 = self.de_block3(c4)
        d3 = F.interpolate(d3, size=e2.size()[2:], mode='bilinear', align_corners=True)
        c3 = torch.cat((d3,e2),1)

        d2 = self.de_block2(c3)
        d2 = F.interpolate(d2, size=e1.size()[2:], mode='bilinear', align_corners=True)
        c2 =torch.cat((d2,e1),1)

        d1 = self.de_block1(c2)
        d1 = F.interpolate(d1, size=e0.size()[2:], mode='bilinear', align_corners=True)
        c1 = torch.cat((d1,e0),1)
        y = self.de_block0(c1)
        
        return y
//  [python]
def track_camera(self, rgb_data: torch.Tensor, depth_data: torch.Tensor,
                     calib: FrameIntrinsic, set_pose: Isometry = None):
        """
        :param rgb_data:    (H, W, 3)       float32
        :param depth_data:  (H, W)          float32
        :param calib:       FrameIntrinsic
        :param set_pose:    Force set pose.
        :return: a pose.
        """
        cur_intensity = torch.mean(rgb_data, dim=-1)
        cur_depth = depth_data
        cur_intensity, cur_depth, cur_dIdxy = self._make_image_pyramid(cur_intensity, cur_depth)
        cur_rgb = rgb_data.permute(2, 0, 1)       # (3, H, W)

        # Process to point cloud.
        pc_scale = self.sdf_args.subsample
        pc_data = t.nn.functional.interpolate(cur_depth[0].unsqueeze(0).unsqueeze(0),
                                                  scale_factor=pc_scale, mode='nearest',
                                                  recompute_scale_factor=False).squeeze(0).squeeze(0)
        cur_rgb = t.nn.functional.interpolate(cur_rgb.unsqueeze(0), scale_factor=pc_scale, mode='bilinear', recompute_scale_factor=False).squeeze(0)
        pc_data = unproject_depth(pc_data, calib.fx * pc_scale, calib.fy * pc_scale,
                                  calib.cx * pc_scale, calib.cy * pc_scale)
        pc_data = torch.cat([pc_data, torch.zeros((pc_data.size(0), pc_data.size(1), 1), device=pc_data.device)], dim=-1)
        pc_data = pc_data.reshape(-1, 4)
        cur_rgb = cur_rgb.permute(1, 2, 0)      # (W, H, 3)
        cur_rgb = cur_rgb.reshape(-1, 3)

        nan_mask = ~torch.isnan(pc_data[..., 0])
        pc_data = pc_data[nan_mask]
        cur_rgb = cur_rgb[nan_mask]

        with torch.cuda.device(self.map.device):
            pc_data_valid_mask = remove_radius_outlier(pc_data, 16, 0.05)
            pc_data = pc_data[pc_data_valid_mask]
            cur_rgb = cur_rgb[pc_data_valid_mask]
            normal_data = estimate_normals(pc_data, 16, 0.1, [0.0, 0.0, 0.0])
            normal_valid_mask = ~torch.isnan(normal_data[..., 0])
            normal_data = normal_data[normal_valid_mask]
            cur_rgb = cur_rgb[normal_valid_mask]
            pc_data = pc_data[normal_valid_mask, :3]

        self.last_colored_pcd = [pc_data, cur_rgb]
        pc_data, normal_data = point_box_filter(pc_data, normal_data, 0.02)
        self.last_processed_pc = [pc_data, normal_data]

        if set_pose is not None:
            final_pose = set_pose
        else:
            assert len(self.all_pd_pose) > 0
            lspeed = Isometry()
            final_pose = self.gauss_newton(self.all_pd_pose[-1].dot(lspeed), cur_intensity, cur_depth, cur_dIdxy, pc_data, calib)

        self.last_intensity = cur_intensity
        self.last_depth = cur_depth
        self.all_pd_pose.append(final_pose)
        return final_pose
//  [python]
def track_camera(self, rgb_data: torch.Tensor, depth_data: torch.Tensor,
                     calib: FrameIntrinsic, set_pose: Isometry = None):
        """
        :param rgb_data:    (H, W, 3)       float32
        :param depth_data:  (H, W)          float32
        :param calib:       FrameIntrinsic
        :param set_pose:    Force set pose.
        :return: a pose.
        """
        cur_intensity = torch.mean(rgb_data, dim=-1)
        cur_depth = depth_data
        cur_intensity, cur_depth, cur_dIdxy = self._make_image_pyramid(cur_intensity, cur_depth)
        cur_rgb = rgb_data.permute(2, 0, 1)       # (3, H, W)

        # Process to point cloud.
        pc_scale = self.sdf_args.subsample
        pc_data = torch.nn.functional.interpolate(cur_depth[0].unsqueeze(0).unsqueeze(0),
                                                  scale_factor=pc_scale, mode='nearest',
                                                  recompute_scale_factor=False).squeeze(0).squeeze(0)
        cur_rgb = torch.nn.functional.interpolate(cur_rgb.unsqueeze(0), scale_factor=pc_scale, mode='bilinear', recompute_scale_factor=False).squeeze(0)
        pc_data = unproject_depth(pc_data, calib.fx * pc_scale, calib.fy * pc_scale,
                                  calib.cx * pc_scale, calib.cy * pc_scale)
        pc_data = torch.cat([pc_data, torch.zeros((pc_data.size(0), pc_data.size(1), 1), device=pc_data.device)], dim=-1)
        pc_data = pc_data.reshape(-1, 4)
        cur_rgb = cur_rgb.permute(1, 2, 0)      # (W, H, 3)
        cur_rgb = cur_rgb.reshape(-1, 3)

        nan_mask = ~torch.isnan(pc_data[..., 0])
        pc_data = pc_data[nan_mask]
        cur_rgb = cur_rgb[nan_mask]

        with torch.cuda.device(self.map.device):
            pc_data_valid_mask = remove_radius_outlier(pc_data, 16, 0.05)
            pc_data = pc_data[pc_data_valid_mask]
            cur_rgb = cur_rgb[pc_data_valid_mask]
            normal_data = estimate_normals(pc_data, 16, 0.1, [0.0, 0.0, 0.0])
            normal_valid_mask = ~torch.isnan(normal_data[..., 0])
            normal_data = normal_data[normal_valid_mask]
            cur_rgb = cur_rgb[normal_valid_mask]
            pc_data = pc_data[normal_valid_mask, :3]

        self.last_colored_pcd = [pc_data, cur_rgb]
        pc_data, normal_data = point_box_filter(pc_data, normal_data, 0.02)
        self.last_processed_pc = [pc_data, normal_data]

        if set_pose is not None:
            final_pose = set_pose
        else:
            assert len(self.all_pd_pose) > 0
            lspeed = Isometry()
            final_pose = self.gauss_newton(self.all_pd_pose[-1].dot(lspeed), cur_intensity, cur_depth, cur_dIdxy, pc_data, calib)

        self.last_intensity = cur_intensity
        self.last_depth = cur_depth
        self.all_pd_pose.append(final_pose)
        return final_pose
//  [python]
def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size
        h0 = h // self.patch_embed.patch_size
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = t.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
            mode='bicubic',
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
//  [python]
def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size
        h0 = h // self.patch_embed.patch_size
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
            mode='bicubic',
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
//  [python]
def _mask_point_forward_test(self, x: Tuple[Tensor], rois: Tensor,
                                 label_preds: Tensor,
                                 mask_preds: Tensor) -> Tensor:
        """Mask refining process with point head in testing.

        Args:
            x (tuple[Tensor]): Feature maps of all scale level.
            rois (Tensor): shape (num_rois, 5).
            label_preds (Tensor): The predication class for each rois.
            mask_preds (Tensor): The predication coarse masks of
                shape (num_rois, num_classes, small_size, small_size).

        Returns:
            Tensor: The refined masks of shape (num_rois, num_classes,
            large_size, large_size).
        """
        refined_mask_pred = mask_preds.clone()
        for subdivision_step in range(self.test_cfg.subdivision_steps):
            refined_mask_pred = t.interpolate(
                refined_mask_pred,
                scale_factor=self.test_cfg.scale_factor,
                mode='bilinear',
                align_corners=False)
            # If `subdivision_num_points` is larger or equal to the
            # resolution of the next step, then we can skip this step
            num_rois, channels, mask_height, mask_width = \
                refined_mask_pred.shape
            if (self.test_cfg.subdivision_num_points >=
                    self.test_cfg.scale_factor**2 * mask_height * mask_width
                    and
                    subdivision_step < self.test_cfg.subdivision_steps - 1):
                continue
            point_indices, rel_roi_points = \
                self.point_head.get_roi_rel_points_test(
                    refined_mask_pred, label_preds, cfg=self.test_cfg)

            fine_grained_point_feats = self._get_fine_grained_point_feats(
                x=x, rois=rois, rel_roi_points=rel_roi_points)
            coarse_point_feats = point_sample(mask_preds, rel_roi_points)
            mask_point_pred = self.point_head(fine_grained_point_feats,
                                              coarse_point_feats)

            point_indices = point_indices.unsqueeze(1).expand(-1, channels, -1)
            refined_mask_pred = refined_mask_pred.reshape(
                num_rois, channels, mask_height * mask_width)
            refined_mask_pred = refined_mask_pred.scatter_(
                2, point_indices, mask_point_pred)
            refined_mask_pred = refined_mask_pred.view(num_rois, channels,
                                                       mask_height, mask_width)

        return refined_mask_pred
//  [python]
def _mask_point_forward_test(self, x: Tuple[Tensor], rois: Tensor,
                                 label_preds: Tensor,
                                 mask_preds: Tensor) -> Tensor:
        """Mask refining process with point head in testing.

        Args:
            x (tuple[Tensor]): Feature maps of all scale level.
            rois (Tensor): shape (num_rois, 5).
            label_preds (Tensor): The predication class for each rois.
            mask_preds (Tensor): The predication coarse masks of
                shape (num_rois, num_classes, small_size, small_size).

        Returns:
            Tensor: The refined masks of shape (num_rois, num_classes,
            large_size, large_size).
        """
        refined_mask_pred = mask_preds.clone()
        for subdivision_step in range(self.test_cfg.subdivision_steps):
            refined_mask_pred = F.interpolate(
                refined_mask_pred,
                scale_factor=self.test_cfg.scale_factor,
                mode='bilinear',
                align_corners=False)
            # If `subdivision_num_points` is larger or equal to the
            # resolution of the next step, then we can skip this step
            num_rois, channels, mask_height, mask_width = \
                refined_mask_pred.shape
            if (self.test_cfg.subdivision_num_points >=
                    self.test_cfg.scale_factor**2 * mask_height * mask_width
                    and
                    subdivision_step < self.test_cfg.subdivision_steps - 1):
                continue
            point_indices, rel_roi_points = \
                self.point_head.get_roi_rel_points_test(
                    refined_mask_pred, label_preds, cfg=self.test_cfg)

            fine_grained_point_feats = self._get_fine_grained_point_feats(
                x=x, rois=rois, rel_roi_points=rel_roi_points)
            coarse_point_feats = point_sample(mask_preds, rel_roi_points)
            mask_point_pred = self.point_head(fine_grained_point_feats,
                                              coarse_point_feats)

            point_indices = point_indices.unsqueeze(1).expand(-1, channels, -1)
            refined_mask_pred = refined_mask_pred.reshape(
                num_rois, channels, mask_height * mask_width)
            refined_mask_pred = refined_mask_pred.scatter_(
                2, point_indices, mask_point_pred)
            refined_mask_pred = refined_mask_pred.view(num_rois, channels,
                                                       mask_height, mask_width)

        return refined_mask_pred
//  [python]
def forward(self, inputs):
        """Forward function."""
        feats = [
            adapt_conv(inputs[i + self.start_level])
            for i, adapt_conv in enumerate(self.adapt_convs)
        ]

        for (i, module_name) in enumerate(self.fpn):
            idx_1, idx_2 = int(module_name[1]), int(module_name[2])
            res = self.fpn[module_name](feats[idx_1], feats[idx_2])
            feats.append(res)

        ret = []
        for (idx, input_idx) in zip([9, 8, 7], [1, 2, 3]):  # add P3, P4, P5
            feats1, feats2 = feats[idx], feats[5]
            feats2_resize = t.interpolate(
                feats2,
                size=feats1.size()[2:],
                mode='bilinear',
                align_corners=False)

            feats_sum = feats1 + feats2_resize
            ret.append(
                t.interpolate(
                    feats_sum,
                    size=inputs[input_idx].size()[2:],
                    mode='bilinear',
                    align_corners=False))

        for submodule in self.extra_downsamples:
            ret.append(submodule(ret[-1]))

        return tuple(ret)
//  [python]
def forward(self, inputs):
        """Forward function."""
        feats = [
            adapt_conv(inputs[i + self.start_level])
            for i, adapt_conv in enumerate(self.adapt_convs)
        ]

        for (i, module_name) in enumerate(self.fpn):
            idx_1, idx_2 = int(module_name[1]), int(module_name[2])
            res = self.fpn[module_name](feats[idx_1], feats[idx_2])
            feats.append(res)

        ret = []
        for (idx, input_idx) in zip([9, 8, 7], [1, 2, 3]):  # add P3, P4, P5
            feats1, feats2 = feats[idx], feats[5]
            feats2_resize = F.interpolate(
                feats2,
                size=feats1.size()[2:],
                mode='bilinear',
                align_corners=False)

            feats_sum = feats1 + feats2_resize
            ret.append(
                F.interpolate(
                    feats_sum,
                    size=inputs[input_idx].size()[2:],
                    mode='bilinear',
                    align_corners=False))

        for submodule in self.extra_downsamples:
            ret.append(submodule(ret[-1]))

        return tuple(ret)
//  [python]
def paint(sampler, image, prompt, t_enc, seed, scale, num_samples=1, callback=None,
          do_full_sample=False):
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = sampler.model
    seed_everything(seed)

    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")
    wm = "SDV2"
    wm_encoder = WatermarkEncoder()
    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))

    with torch.no_grad(),\
            torch.autocast("cuda"):
        batch = make_batch_sd(image, txt=prompt, device=device, num_samples=num_samples)
        z = model.get_first_stage_encoding(model.encode_first_stage(batch[model.first_stage_key]))  # move to latent space
        c = model.cond_stage_model.encode(batch["txt"])
        c_cat = list()
        for ck in model.concat_keys:
            cc = batch[ck]
            cc = model.depth_model(cc)
            depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
                                                                                           keepdim=True)
            display_depth = (cc - depth_min) / (depth_max - depth_min)
            st.image(Image.fromarray((display_depth[0, 0, ...].cpu().numpy() * 255.).astype(np.uint8)))
            cc = t.interpolate(
                cc,
                size=z.shape[2:],
                mode="bicubic",
                align_corners=False,
            )
            depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
                                                                                           keepdim=True)
            cc = 2. * (cc - depth_min) / (depth_max - depth_min) - 1.
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        # cond
        cond = {"c_concat": [c_cat], "c_crossattn": [c]}

        # uncond cond
        uc_cross = model.get_unconditional_conditioning(num_samples, "")
        uc_full = {"c_concat": [c_cat], "c_crossattn": [uc_cross]}
        if not do_full_sample:
            # encode (scaled latent)
            z_enc = sampler.stochastic_encode(z, torch.tensor([t_enc] * num_samples).to(model.device))
        else:
            z_enc = torch.randn_like(z)
        # decode it
        samples = sampler.decode(z_enc, cond, t_enc, unconditional_guidance_scale=scale,
                                 unconditional_conditioning=uc_full, callback=callback)
        x_samples_ddim = model.decode_first_stage(samples)
        result = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
        result = result.cpu().numpy().transpose(0, 2, 3, 1) * 255
    return [put_watermark(Image.fromarray(img.astype(np.uint8)), wm_encoder) for img in result]
//  [python]
def paint(sampler, image, prompt, t_enc, seed, scale, num_samples=1, callback=None,
          do_full_sample=False):
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = sampler.model
    seed_everything(seed)

    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")
    wm = "SDV2"
    wm_encoder = WatermarkEncoder()
    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))

    with torch.no_grad(),\
            torch.autocast("cuda"):
        batch = make_batch_sd(image, txt=prompt, device=device, num_samples=num_samples)
        z = model.get_first_stage_encoding(model.encode_first_stage(batch[model.first_stage_key]))  # move to latent space
        c = model.cond_stage_model.encode(batch["txt"])
        c_cat = list()
        for ck in model.concat_keys:
            cc = batch[ck]
            cc = model.depth_model(cc)
            depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
                                                                                           keepdim=True)
            display_depth = (cc - depth_min) / (depth_max - depth_min)
            st.image(Image.fromarray((display_depth[0, 0, ...].cpu().numpy() * 255.).astype(np.uint8)))
            cc = torch.nn.functional.interpolate(
                cc,
                size=z.shape[2:],
                mode="bicubic",
                align_corners=False,
            )
            depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
                                                                                           keepdim=True)
            cc = 2. * (cc - depth_min) / (depth_max - depth_min) - 1.
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        # cond
        cond = {"c_concat": [c_cat], "c_crossattn": [c]}

        # uncond cond
        uc_cross = model.get_unconditional_conditioning(num_samples, "")
        uc_full = {"c_concat": [c_cat], "c_crossattn": [uc_cross]}
        if not do_full_sample:
            # encode (scaled latent)
            z_enc = sampler.stochastic_encode(z, torch.tensor([t_enc] * num_samples).to(model.device))
        else:
            z_enc = torch.randn_like(z)
        # decode it
        samples = sampler.decode(z_enc, cond, t_enc, unconditional_guidance_scale=scale,
                                 unconditional_conditioning=uc_full, callback=callback)
        x_samples_ddim = model.decode_first_stage(samples)
        result = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
        result = result.cpu().numpy().transpose(0, 2, 3, 1) * 255
    return [put_watermark(Image.fromarray(img.astype(np.uint8)), wm_encoder) for img in result]
//  [python]
def encode(self, input):
        h, w = input.size()[-2:]
        sh, sw = h//2**self.scale_ratio, w//2**self.scale_ratio
        x = F.upsample(input, size=(sh, sw))
        return self.fc(x)
//  [python]
def encode(self, input):
        h, w = input.size()[-2:]
        sh, sw = h//2**self.scale_ratio, w//2**self.scale_ratio
        x = F.interpolate(input, size=(sh, sw))
        return self.fc(x)
//  [python]
def _load_pretrained(model: TinyViT, url: str) -> TinyViT:
    model_state_dict = model.state_dict()
    state_dict = torch.hub.load_state_dict_from_url(url)

    # official checkpoint has "model" key
    if "model" in state_dict:
        state_dict = state_dict["model"]

    # https://github.com/microsoft/Cream/blob/8dc38822b99fff8c262c585a32a4f09ac504d693/TinyViT/utils.py#L163
    # bicubic interpolate attention biases
    ab_keys = [k for k in state_dict.keys() if "attention_biases" in k]
    for k in ab_keys:
        n_heads1, L1 = state_dict[k].shape
        n_heads2, L2 = model_state_dict[k].shape
        KORNIA_CHECK(n_heads1 == n_heads2, f"Fail to load {k}. Pre-trained checkpoint should have num_heads={n_heads1}")

        if L1 != L2:
            S1 = int(L1**0.5)
            S2 = int(L2**0.5)
            attention_biases = state_dict[k].view(1, n_heads1, S1, S1)
            attention_biases = t.interpolate(attention_biases, size=(S2, S2), mode="bicubic")
            state_dict[k] = attention_biases.view(n_heads2, L2)

    if state_dict["head.weight"].shape[0] != model.head.out_features:
        msg = "Number of classes does not match pre-trained checkpoint's. Resetting classification head to zeros"
        warnings.warn(msg)
        state_dict["head.weight"] = torch.zeros_like(model.head.weight)
        state_dict["head.bias"] = torch.zeros_like(model.head.bias)

    model.load_state_dict(state_dict)
    return model
//  [python]
def _load_pretrained(model: TinyViT, url: str) -> TinyViT:
    model_state_dict = model.state_dict()
    state_dict = torch.hub.load_state_dict_from_url(url)

    # official checkpoint has "model" key
    if "model" in state_dict:
        state_dict = state_dict["model"]

    # https://github.com/microsoft/Cream/blob/8dc38822b99fff8c262c585a32a4f09ac504d693/TinyViT/utils.py#L163
    # bicubic interpolate attention biases
    ab_keys = [k for k in state_dict.keys() if "attention_biases" in k]
    for k in ab_keys:
        n_heads1, L1 = state_dict[k].shape
        n_heads2, L2 = model_state_dict[k].shape
        KORNIA_CHECK(n_heads1 == n_heads2, f"Fail to load {k}. Pre-trained checkpoint should have num_heads={n_heads1}")

        if L1 != L2:
            S1 = int(L1**0.5)
            S2 = int(L2**0.5)
            attention_biases = state_dict[k].view(1, n_heads1, S1, S1)
            attention_biases = F.interpolate(attention_biases, size=(S2, S2), mode="bicubic")
            state_dict[k] = attention_biases.view(n_heads2, L2)

    if state_dict["head.weight"].shape[0] != model.head.out_features:
        msg = "Number of classes does not match pre-trained checkpoint's. Resetting classification head to zeros"
        warnings.warn(msg)
        state_dict["head.weight"] = torch.zeros_like(model.head.weight)
        state_dict["head.bias"] = torch.zeros_like(model.head.bias)

    model.load_state_dict(state_dict)
    return model