Function frame_transform

streambridge/utils.py:209–246

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

    return Compose(transforms)

def frame_transform(
        image_size: Union[int, Tuple[int, 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, int):
        resize_size = (image_size, image_size)
        crop_size = (image_size, image_size)
    elif isinstance(image_size, (list, tuple)) and len(image_size) == 2:
        resize_size = (image_size[0], image_size[1])
        crop_size = (image_size[0], image_size[1])
    else:
        raise ValueError("image_size must be an int or a tuple of two ints.")

    normalize = Normalize(mean=mean, std=std)
    
    transforms = [
        ToPILImage(),
        Resize(resize_size, interpolation=InterpolationMode.BICUBIC),
        CenterCrop(resize_size),
        RandomHorizontalFlip() if random_flip else Lambda(lambda x: x),
    ]
    transforms.extend([
        _convert_to_rgb,
        ToTensor(),
        normalize,
    ])
    return Compose(transforms)

def expand2square(pil_img, background_color=tuple(int(x*255) for x in OPENAI_DATASET_MEAN)):
    width, height = pil_img.size