(tensor_list: List[Tensor])
| 46 | |
| 47 | |
| 48 | def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): |
| 49 | # TODO make this more general |
| 50 | if tensor_list[0].ndim == 3: |
| 51 | if torchvision._is_tracing(): |
| 52 | # nested_tensor_from_tensor_list() does not export well to ONNX |
| 53 | # call _onnx_nested_tensor_from_tensor_list() instead |
| 54 | return _onnx_nested_tensor_from_tensor_list(tensor_list) |
| 55 | |
| 56 | # TODO make it support different-sized images |
| 57 | 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_size |
| 60 | b, c, h, w = batch_shape |
| 61 | dtype = tensor_list[0].dtype |
| 62 | device = tensor_list[0].device |
| 63 | 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]] = False |
| 68 | else: |
| 69 | raise ValueError("not supported") |
| 70 | return NestedTensor(tensor, mask) |
| 71 | |
| 72 | |
| 73 | # _onnx_nested_tensor_from_tensor_list() is an implementation of |
nothing calls this directly
no test coverage detected