Voxel51/openworld-sam

1# Copyright (c) Facebook, Inc. and its affiliates.2# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/util/misc.py3"""4Misc functions, including distributed helpers.5 6Mostly copy-paste from torchvision references.7"""8from typing import List, Optional9 10import torch11import torch.distributed as dist12import torchvision13from torch import Tensor14 15 16def _max_by_axis(the_list):17    # type: (List[List[int]]) -> List[int]18    maxes = the_list[0]19    for sublist in the_list[1:]:20        for index, item in enumerate(sublist):21            maxes[index] = max(maxes[index], item)22    return maxes23 24 25class NestedTensor(object):26    def __init__(self, tensors, mask: Optional[Tensor]):27        self.tensors = tensors28        self.mask = mask29 30    def to(self, device):31        # type: (Device) -> NestedTensor # noqa32        cast_tensor = self.tensors.to(device)33        mask = self.mask34        if mask is not None:35            assert mask is not None36            cast_mask = mask.to(device)37        else:38            cast_mask = None39        return NestedTensor(cast_tensor, cast_mask)40 41    def decompose(self):42        return self.tensors, self.mask43 44    def __repr__(self):45        return str(self.tensors)46 47 48def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):49    # TODO make this more general50    if tensor_list[0].ndim == 3:51        if torchvision._is_tracing():52            # nested_tensor_from_tensor_list() does not export well to ONNX53            # call _onnx_nested_tensor_from_tensor_list() instead54            return _onnx_nested_tensor_from_tensor_list(tensor_list)55 56        # TODO make it support different-sized images57        max_size = _max_by_axis([list(img.shape) for img in tensor_list])58        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))59        batch_shape = [len(tensor_list)] + max_size60        b, c, h, w = batch_shape61        dtype = tensor_list[0].dtype62        device = tensor_list[0].device63        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)64        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)65        for img, pad_img, m in zip(tensor_list, tensor, mask):66            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)67            m[: img.shape[1], : img.shape[2]] = False68    else:69        raise ValueError("not supported")70    return NestedTensor(tensor, mask)71 72 73# _onnx_nested_tensor_from_tensor_list() is an implementation of74# nested_tensor_from_tensor_list() that is supported by ONNX tracing.75@torch.jit.unused76def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:77    max_size = []78    for i in range(tensor_list[0].dim()):79        max_size_i = torch.max(80            torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)81        ).to(torch.int64)82        max_size.append(max_size_i)83    max_size = tuple(max_size)84 85    # work around for86    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)87    # m[: img.shape[1], :img.shape[2]] = False88    # which is not yet supported in onnx89    padded_imgs = []90    padded_masks = []91    for img in tensor_list:92        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]93        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))94        padded_imgs.append(padded_img)95 96        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)97        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)98        padded_masks.append(padded_mask.to(torch.bool))99 100    tensor = torch.stack(padded_imgs)101    mask = torch.stack(padded_masks)102 103    return NestedTensor(tensor, mask=mask)104 105 106def is_dist_avail_and_initialized():107    if not dist.is_available():108        return False109    if not dist.is_initialized():110        return False111    return True112