Spaces:
Sleeping
Sleeping
File size: 6,020 Bytes
34f3bc9 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | import torch
import torch.nn.functional as F
import math
def make_bool_mask(*dims: int) -> torch.BoolTensor:
"""Make a boolean mask for the given dimensions.
Example:
_make_bool_mask(2, -2, 2) == (True, True, False, False, True, True)
_make_bool_mask(2, 0, 2) == (True, True, True, True)
Args:
dims: The dimensions to make the mask for.
Returns:
A tensor of booleans.
"""
lens = [abs(d) for d in dims]
flags = [d > 0 for d in dims]
parts = [torch.full((l,), f, dtype=torch.bool) for l, f in zip(lens, flags)]
return torch.cat(parts) if parts else torch.empty(0, dtype=torch.bool)
def resize_with_pad(
image: torch.Tensor,
target_h: int,
target_w: int,
mode: str = "bilinear",
) -> torch.Tensor:
"""
Resize a tensor to target size while keeping aspect ratio.
Supports both [3, H, W] and [B, 3, H, W] inputs.
Vectorized over the batch dimension (no Python loop).
Args:
image: torch.Tensor, shape [3, H, W] or [B, 3, H, W], values in [0, 1].
target_h: target height.
target_w: target width.
mode: interpolation mode ('bilinear', 'nearest', etc.).
Returns:
torch.Tensor of shape [3, target_h, target_w] or [B, 3, target_h, target_w].
"""
if image.ndim == 3:
# [3, H, W] -> [1, 3, H, W]
image = image.unsqueeze(0)
squeeze_back = True
elif image.ndim == 4:
squeeze_back = False
else:
raise ValueError(f"Expected tensor of shape [3,H,W] or [B,3,H,W], got {image.shape}")
B, C, H, W = image.shape
if H == target_h and W == target_w:
return image.squeeze(0) if squeeze_back else image
scale = min(target_h / H, target_w / W)
new_h, new_w = int(round(H * scale)), int(round(W * scale))
image_resized = F.interpolate(
image,
size=(new_h, new_w),
mode=mode,
align_corners=False if mode == "bilinear" else None,
) # [B, 3, new_h, new_w]
pad_top = (target_h - new_h) // 2
pad_bottom = target_h - new_h - pad_top
pad_left = (target_w - new_w) // 2
pad_right = target_w - new_w - pad_left
image_padded = F.pad(image_resized, (pad_left, pad_right, pad_top, pad_bottom), value=0.0)
image_padded = image_padded.clamp(0.0, 1.0)
# sanity check
assert image_padded.shape[-2:] == (target_h, target_w)
return image_padded.squeeze(0) if squeeze_back else image_padded
def resize_center_crop(
image: torch.Tensor,
target_h: int,
target_w: int,
mode: str = "bilinear",
) -> torch.Tensor:
"""
Resize an image or a batch of images so that the shortest side is scaled
to at least the target size, then apply a center crop to exactly
[C, target_h, target_w] or [B, C, target_h, target_w].
Supports:
- [C, H, W]
- [B, C, H, W]
and performs vectorized operations over the batch dimension.
Args:
image: Input tensor of shape [C,H,W] or [B,C,H,W].
target_h: Target height after cropping.
target_w: Target width after cropping.
mode: Interpolation mode for resizing.
Returns:
Tensor of shape [C, target_h, target_w] or [B, C, target_h, target_w].
"""
# ------------------------------------------------------
# Normalize input shape to batch mode
# ------------------------------------------------------
if image.ndim == 3:
# Convert [C, H, W] → [1, C, H, W]
assert image.shape[0] in (1, 3), f"Expected [C,H,W], got {image.shape}"
image = image.unsqueeze(0)
squeeze_back = True
elif image.ndim == 4:
squeeze_back = False
else:
raise ValueError(f"Expected [C,H,W] or [B,C,H,W], got {image.shape}")
B, C, H, W = image.shape
# If the image is already the correct size, return early
if (H, W) == (target_h, target_w):
return image.squeeze(0) if squeeze_back else image
# ------------------------------------------------------
# Compute scale factor so the shortest side >= target size
# ------------------------------------------------------
# Use max() to ensure both resized dimensions are >= target dimensions.
scale = max(target_h / H, target_w / W)
# Ceil prevents rounding down to a size smaller than required
new_h = int(math.ceil(H * scale))
new_w = int(math.ceil(W * scale))
# ------------------------------------------------------
# Resize the batch in a single vectorized operation
# ------------------------------------------------------
align_corners = False if mode in ("bilinear", "bicubic") else None
x = F.interpolate(
image,
size=(new_h, new_w),
mode=mode,
align_corners=align_corners,
) # [B, C, new_h, new_w]
# ------------------------------------------------------
# Center crop to the target size
# ------------------------------------------------------
top = max((new_h - target_h) // 2, 0)
left = max((new_w - target_w) // 2, 0)
x = x[:, :, top : top + target_h, left : left + target_w]
# ------------------------------------------------------
# (Rare) If rounding left us with slightly smaller output,
# add minimal padding to reach the exact target size
# ------------------------------------------------------
cur_h, cur_w = x.shape[-2:]
pad_h = target_h - cur_h
pad_w = target_w - cur_w
if pad_h > 0 or pad_w > 0:
# F.pad padding order: (left, right, top, bottom)
x = F.pad(
x,
(0, max(pad_w, 0), 0, max(pad_h, 0)), # Only pad right/bottom
)
# Clamp values to valid range
x = x.clamp(0.0, 1.0)
# Sanity check
assert x.shape[-2:] == (target_h, target_w)
# ------------------------------------------------------
# Return a squeezed tensor if input was a single image
# ------------------------------------------------------
return x.squeeze(0) if squeeze_back else x
|