Spaces:
Running
Running
File size: 2,626 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import os
import torch
import torchvision.transforms as transforms
from functools import lru_cache
@lru_cache(maxsize=None)
def meshgrid(B, H, W, dtype, device, normalized=False):
"""
Create mesh-grid given batch size, height and width dimensions. From https://github.com/TRI-ML/KP2D.
Parameters
----------
B: int
Batch size
H: int
Grid Height
W: int
Batch size
dtype: torch.dtype
Tensor dtype
device: str
Tensor device
normalized: bool
Normalized image coordinates or integer-grid.
Returns
-------
xs: torch.Tensor
Batched mesh-grid x-coordinates (BHW).
ys: torch.Tensor
Batched mesh-grid y-coordinates (BHW).
"""
if normalized:
xs = torch.linspace(-1, 1, W, device=device, dtype=dtype)
ys = torch.linspace(-1, 1, H, device=device, dtype=dtype)
else:
xs = torch.linspace(0, W - 1, W, device=device, dtype=dtype)
ys = torch.linspace(0, H - 1, H, device=device, dtype=dtype)
ys, xs = torch.meshgrid([ys, xs])
return xs.repeat([B, 1, 1]), ys.repeat([B, 1, 1])
@lru_cache(maxsize=None)
def image_grid(B, H, W, dtype, device, ones=True, normalized=False):
"""
Create an image mesh grid with shape B3HW given image shape BHW. From https://github.com/TRI-ML/KP2D.
Parameters
----------
B: int
Batch size
H: int
Grid Height
W: int
Batch size
dtype: str
Tensor dtype
device: str
Tensor device
ones : bool
Use (x, y, 1) coordinates
normalized: bool
Normalized image coordinates or integer-grid.
Returns
-------
grid: torch.Tensor
Mesh-grid for the corresponding image shape (B3HW)
"""
xs, ys = meshgrid(B, H, W, dtype, device, normalized=normalized)
coords = [xs, ys]
if ones:
coords.append(torch.ones_like(xs)) # BHW
grid = torch.stack(coords, dim=1) # B3HW
return grid
def to_tensor_sample(sample, tensor_type="torch.FloatTensor"):
"""
Casts the keys of sample to tensors. From https://github.com/TRI-ML/KP2D.
Parameters
----------
sample : dict
Input sample
tensor_type : str
Type of tensor we are casting to
Returns
-------
sample : dict
Sample with keys cast as tensors
"""
transform = transforms.ToTensor()
sample["image"] = transform(sample["image"]).type(tensor_type)
return sample
def prepare_dirs(config):
for path in [config.ckpt_dir]:
if not os.path.exists(path):
os.makedirs(path)
|