76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from torchvision.transforms import v2
|
|
|
|
|
|
class ResizeAndPadBatch:
|
|
"""
|
|
Vectorized resize + zero pad for batched tensors.
|
|
|
|
Args:
|
|
imgs (Tensor): shape [B, C, H, W], pixel values in [0,1] or [0,255].
|
|
target_size (tuple): (target_h, target_w) final image shape.
|
|
|
|
Returns:
|
|
Tensor: shape [B, C, target_h, target_w]
|
|
"""
|
|
|
|
def __init__(self, target_size=(640, 640), fill_value=0):
|
|
self.target_size = target_size
|
|
self.fill_value = fill_value
|
|
self.transform_stats = None
|
|
|
|
def __call__(self, imgs: torch.Tensor):
|
|
|
|
b, c, h, w = imgs.shape
|
|
target_h, target_w = self.target_size
|
|
|
|
# Compute scaling factors: preserve aspect ratio by same scale factor per image
|
|
scale_factors = float(
|
|
torch.minimum(torch.tensor(target_h / h), torch.tensor(target_w / w))
|
|
)
|
|
|
|
# New intermediate size
|
|
new_h = int(h * scale_factors)
|
|
new_w = int(w * scale_factors)
|
|
|
|
# Resize with bilinear interpolation
|
|
resized = F.interpolate(
|
|
imgs, size=(new_h, new_w), mode="bilinear", align_corners=False
|
|
)
|
|
|
|
# Calculate padding amounts (left, right, top, bottom)
|
|
pad_h = target_h - new_h
|
|
pad_w = target_w - new_w
|
|
pad_top = pad_h // 2
|
|
pad_bottom = pad_h - pad_top
|
|
pad_left = pad_w // 2
|
|
pad_right = pad_w - pad_left
|
|
|
|
# Apply padding (pad order in F.pad is [left, right, top, bottom])
|
|
padded = F.pad(
|
|
resized,
|
|
(pad_left, pad_right, pad_top, pad_bottom),
|
|
mode="constant",
|
|
value=self.fill_value,
|
|
)
|
|
|
|
self.transform_stats = {
|
|
"scale_factor": scale_factors,
|
|
'hw': [h,w],
|
|
"new_hw": [new_h, new_w],
|
|
"pad_TBLR": [pad_top,pad_bottom, pad_left, pad_right]
|
|
}
|
|
|
|
return padded
|
|
|
|
|
|
clip_transforms = v2.Compose(
|
|
[
|
|
v2.Resize(size=(512, 512)),
|
|
v2.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
|
|
]
|
|
)
|
|
|
|
det_transforms = v2.Compose([ResizeAndPadBatch(target_size=(640, 640))])
|