Alternate depth normalization about vit-lens HOT 4 OPEN

import torch
import torchvision.transforms.functional as F
from torchvision.transforms import RandomResizedCrop

class RandomResizedCropAndScale(RandomResizedCrop):
    """
    Crop a random portion of image and resize it to a given size. Scale it by the sqrt of the ratio in areas between the new image and the crop size in the original image, but only apply this scaling to the final channel.
    """

    def __init__(self, size, *args, **kwargs):
        super().__init__(size, *args, **kwargs)

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Image to be cropped and resized.

        Returns:
            Tensor: Randomly cropped, resized, and scaled image.
        """
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
        cropped_and_resized_img = F.resized_crop(img, i, j, h, w, self.size, self.interpolation, antialias=self.antialias)

        # Convert the cropped and resized image to a tensor if it's not already
        if not isinstance(cropped_and_resized_img, torch.Tensor):
            cropped_and_resized_img = F.to_tensor(cropped_and_resized_img)

        _, height, width = F.get_dimensions(cropped_and_resized_img)
        
        scale_factor = torch.sqrt((height * width) / (h * w))
        
        scaled_img = cropped_and_resized_img.clone()  # Clone to avoid in-place modification issues
        scaled_img[-1, :, :] *= scale_factor  # This applies scale_factor to the last channel


        return scaled_img

    def __call__(self, depth):
        # here depth refers to disparity, in torch savefile format
        # note use ToTensor to scale image to [0,1] first
        img = torch.randn((3, 224, 224))

        if depth.ndim == 2:
            depth = depth.unsqueeze(0)
			
	scale = 224/depth.shape[0]

        rgbd = torch.cat([img, depth * scale], dim=0)
        transform_rgbd = self.rgbd_transform(rgbd)
        img = transform_rgbd[0:3, ...]
        depth = transform_rgbd[3:4, ...]
        
        return depth