Function image_transform

streambridge/utils.py:175–207

(
        image_size: int,
        rescale_factor: float = 1.0,
        mean: Optional[Tuple[float, ...]] = None,
        std: Optional[Tuple[float, ...]] = None,
        random_flip = False,
)

INTERNVIDEO_STD = (0.229, 0.224, 0.225)

def image_transform(
        image_size: int,
        rescale_factor: float = 1.0,
        mean: Optional[Tuple[float, ...]] = None,
        std: Optional[Tuple[float, ...]] = None,
        random_flip = False,
):
    mean = mean or OPENAI_DATASET_MEAN
    if not isinstance(mean, (list, tuple)):
        mean = (mean,) * 3

    std = std or OPENAI_DATASET_STD
    if not isinstance(std, (list, tuple)):
        std = (std,) * 3
    
    if isinstance(image_size, (list, tuple)) and image_size[0] == image_size[1]:
        # for square size, pass size as int so that Resize() uses aspect preserving shortest edge
        image_size = image_size[0]

    normalize = Normalize(mean=mean, std=std)
    
    transforms = [
        Resize(image_size, interpolation=InterpolationMode.BICUBIC),
        CenterCrop(image_size),
        RandomHorizontalFlip() if random_flip else Lambda(lambda x: x),
    ]

    transforms.extend([
        _convert_to_rgb,
        ToTensor(),
        normalize,
    ])
    return Compose(transforms)

def frame_transform(
        image_size: Union[int, Tuple[int, int]],