Spaces:
Running on Zero
Running on Zero
| """ | |
| copy the depth norm code from marigold | |
| """ | |
| import random | |
| from PIL import Image | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from torchvision import transforms | |
| from torchvision.transforms import functional as F | |
| from torchvision.transforms import InterpolationMode | |
| from loguru import logger | |
| from .transforms import MaxLongEdgeMinShortEdgeResize | |
| from .depth_encoding.power_transform import DepthPowerTransform | |
| from .depth_encoding.multiscale_power_transform import DepthMultiScalePowerTransform | |
| from .depth_encoding.disparity_transform import DepthDisparityTransform | |
| from .depth_encoding.multiscale_disparity_transform import DepthMultiScaleDisparityTransform | |
| from .depth_encoding.scale_shift_transform import ScaleShiftDepthNormalizer | |
| class DepthTransform: | |
| def __init__( | |
| self, | |
| max_image_size, | |
| min_image_size, | |
| image_stride, | |
| max_pixels=14*14*9*1024, | |
| kitti_bm_crop=False, | |
| **kwargs | |
| ): | |
| self.stride = image_stride | |
| self.resize_transform = MaxLongEdgeMinShortEdgeResize( | |
| max_size=max_image_size, | |
| min_size=min_image_size, | |
| stride=image_stride, | |
| max_pixels=max_pixels, | |
| ) | |
| self.to_tensor_transform = transforms.ToTensor() | |
| self.kitti_bm_crop = kitti_bm_crop | |
| def __call__(self, img, img_num=1): | |
| img = self.resize_transform(img, img_num=img_num) | |
| if isinstance(img, Image.Image) or isinstance(img, np.ndarray): | |
| img = self.to_tensor_transform(img) | |
| return img | |
| def kitti_benchmark_crop_tensor(input_img): | |
| """ | |
| Crop images to KITTI benchmark size | |
| Args: | |
| `input_img` (torch.Tensor): Input image to be cropped. | |
| Returns: | |
| torch.Tensor:Cropped image. | |
| """ | |
| KB_CROP_HEIGHT = 352 | |
| KB_CROP_WIDTH = 1216 | |
| height, width = input_img.shape[-2:] | |
| top_margin = int(height - KB_CROP_HEIGHT) | |
| left_margin = int((width - KB_CROP_WIDTH) / 2) | |
| if 2 == len(input_img.shape): | |
| out = input_img[ | |
| top_margin : top_margin + KB_CROP_HEIGHT, | |
| left_margin : left_margin + KB_CROP_WIDTH, | |
| ] | |
| elif 3 == len(input_img.shape): | |
| out = input_img[ | |
| :, | |
| top_margin : top_margin + KB_CROP_HEIGHT, | |
| left_margin : left_margin + KB_CROP_WIDTH, | |
| ] | |
| return out | |
| def kitti_benchmark_crop_pil(input_img): | |
| """ | |
| Crop images to KITTI benchmark size | |
| Args: | |
| `input_img` (PIL.Image): Input image to be cropped. | |
| Returns: | |
| PIL.Image: Cropped image. | |
| """ | |
| KB_CROP_HEIGHT = 352 | |
| KB_CROP_WIDTH = 1216 | |
| width, height = input_img.size | |
| top_margin = height - KB_CROP_HEIGHT | |
| left_margin = (width - KB_CROP_WIDTH) // 2 | |
| right_margin = left_margin + KB_CROP_WIDTH | |
| bottom_margin = top_margin + KB_CROP_HEIGHT | |
| out = input_img.crop((left_margin, top_margin, right_margin, bottom_margin)) | |
| return out | |
| class DepthTransformBase: | |
| def __init__( | |
| self, | |
| max_image_size, | |
| min_image_size, | |
| image_stride, | |
| max_pixels=14*14*9*1024, | |
| kitti_bm_crop=False, | |
| normalize=None, | |
| stack=False, | |
| **kwargs | |
| ): | |
| self.stride = image_stride | |
| self.resize_transform = MaxLongEdgeMinShortEdgeResize( | |
| max_size=max_image_size, | |
| min_size=min_image_size, | |
| stride=image_stride, | |
| max_pixels=max_pixels, | |
| ) | |
| self.to_tensor_transform = transforms.ToTensor() | |
| self.kitti_bm_crop = kitti_bm_crop | |
| self.normalize = normalize | |
| self.stack = stack | |
| def __call__(self, img, img_num=1, valid_mask=None): | |
| """ | |
| image: shape=shape=(1, h, w), tensor | |
| """ | |
| # convert to tensor | |
| if isinstance(img, Image.Image): | |
| img = np.asarray(img) # float 32 format | |
| if isinstance(img, np.ndarray): | |
| img = self.to_tensor_transform(img) | |
| # resize | |
| img = self.resize_transform(img, img_num=img_num) | |
| # normalize | |
| if self.normalize is not None: | |
| img = self.normalize(img, valid_mask=valid_mask) | |
| # copy | |
| if self.stack: | |
| img = img.repeat(3, 1, 1) | |
| return img | |
| def kitti_benchmark_crop_tensor(input_img): | |
| """ | |
| Crop images to KITTI benchmark size | |
| Args: | |
| `input_img` (torch.Tensor): Input image to be cropped. | |
| Returns: | |
| torch.Tensor:Cropped image. | |
| """ | |
| KB_CROP_HEIGHT = 352 | |
| KB_CROP_WIDTH = 1216 | |
| height, width = input_img.shape[-2:] | |
| top_margin = int(height - KB_CROP_HEIGHT) | |
| left_margin = int((width - KB_CROP_WIDTH) / 2) | |
| if 2 == len(input_img.shape): | |
| out = input_img[ | |
| top_margin : top_margin + KB_CROP_HEIGHT, | |
| left_margin : left_margin + KB_CROP_WIDTH, | |
| ] | |
| elif 3 == len(input_img.shape): | |
| out = input_img[ | |
| :, | |
| top_margin : top_margin + KB_CROP_HEIGHT, | |
| left_margin : left_margin + KB_CROP_WIDTH, | |
| ] | |
| return out | |
| def kitti_benchmark_crop_pil(input_img): | |
| """ | |
| Crop images to KITTI benchmark size | |
| Args: | |
| `input_img` (PIL.Image): Input image to be cropped. | |
| Returns: | |
| PIL.Image: Cropped image. | |
| """ | |
| KB_CROP_HEIGHT = 352 | |
| KB_CROP_WIDTH = 1216 | |
| width, height = input_img.size | |
| top_margin = height - KB_CROP_HEIGHT | |
| left_margin = (width - KB_CROP_WIDTH) // 2 | |
| right_margin = left_margin + KB_CROP_WIDTH | |
| bottom_margin = top_margin + KB_CROP_HEIGHT | |
| out = input_img.crop((left_margin, top_margin, right_margin, bottom_margin)) | |
| return out | |
| def build_depth_transform(depth_transform_args): | |
| depth_normalization = depth_transform_args.get('depth_normalization', {'type': 'scale_shift'}) | |
| if depth_normalization['type'] == 'scale_shift': | |
| logger.info(f"using scale shift transform") | |
| normalizer = ScaleShiftDepthNormalizer(norm_min=-1.0, norm_max=1.0, min_max_quantile=0.02) | |
| depth_transform = DepthTransformBase( | |
| **depth_transform_args, | |
| normalize=normalizer, | |
| stack=True, | |
| ) | |
| elif depth_normalization['type'] == "power_transform": | |
| logger.info(f"using power transform") | |
| normalizer = DepthPowerTransform( | |
| lam=depth_normalization.get('lam', -3.0), | |
| c=depth_normalization.get('c', 10.0 / 3.0) | |
| ) | |
| depth_transform = DepthTransformBase( | |
| **depth_transform_args, | |
| normalize=normalizer, | |
| stack=True, | |
| ) | |
| elif depth_normalization['type'] == "multiscale_power_transform": | |
| logger.info(f"using power transform") | |
| normalizer = DepthMultiScalePowerTransform( | |
| lam=depth_normalization.get('lam', -3.0), | |
| c=depth_normalization.get('c', [2.0, 6.0, 18.0]) | |
| ) | |
| depth_transform = DepthTransformBase( | |
| **depth_transform_args, | |
| normalize=normalizer, | |
| stack=True, | |
| ) | |
| elif depth_normalization['type'] == "disparity_transform": | |
| logger.info(f"using disparity transform") | |
| normalizer = DepthDisparityTransform( | |
| c=depth_normalization.get('c', 5.0) | |
| ) | |
| depth_transform = DepthTransformBase( | |
| **depth_transform_args, | |
| normalize=normalizer, | |
| stack=True, | |
| ) | |
| elif depth_normalization['type'] == "multiscale_disparity_transform": | |
| logger.info(f"using multiscale disparity transform") | |
| normalizer = DepthMultiScaleDisparityTransform( | |
| c_values=depth_normalization.get('c', [2.0, 8.0, 32.0]) | |
| ) | |
| depth_transform = DepthTransformBase( | |
| **depth_transform_args, | |
| normalize=normalizer, | |
| stack=False, | |
| ) | |
| else: | |
| raise ValueError(f"depth_normalization type {depth_normalization['type']} not supported") | |
| return depth_transform | |
| def decode_depth_wrap(depth_transform: DepthTransformBase, depth_vae_out): | |
| if (depth_transform is None) or isinstance(depth_transform.normalize, ScaleShiftDepthNormalizer): | |
| depth_map = (depth_vae_out * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 0).mean(dim=-1) # (h, w) | |
| elif isinstance(depth_transform.normalize, DepthMultiScaleDisparityTransform) or isinstance(depth_transform.normalize, DepthMultiScalePowerTransform): | |
| # input shape=(3, h, w), output shape=(h, w) | |
| depth_map = depth_transform.normalize.decode(depth_vae_out.clamp(-1.0, 1.0)) | |
| elif isinstance(depth_transform.normalize, DepthDisparityTransform) or isinstance(depth_transform.normalize, DepthPowerTransform): | |
| depth_map = depth_transform.normalize.decode(depth_vae_out.clamp(-1.0, 1.0).mean(dim=0)) | |
| else: | |
| raise ValueError(f"Unknown depth transform type: {depth_transform}") | |
| return depth_map | |