562 lines
18 KiB
Python
562 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import random
|
|
|
|
from packaging import version
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.transforms.functional as functional
|
|
import torchvision.transforms.transforms as transforms
|
|
from scepter.modules.transform.registry import TRANSFORMS
|
|
from scepter.modules.transform.utils import (BACKEND_TORCHVISION,
|
|
INTERPOLATION_STYLE, is_tensor)
|
|
# torchvision.transform._transforms_video is deprecated since torchvision 0.10.0, use transform instead
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from torchvision.version import __version__ as tv_version
|
|
|
|
use_video_transforms = version.parse(tv_version) < version.parse('0.10.0')
|
|
|
|
BACKENDS = (BACKEND_TORCHVISION, )
|
|
|
|
|
|
class VideoTransform(object):
|
|
para_dict = [{
|
|
'INPUT_KEY': {
|
|
'value': 'img',
|
|
'description': 'input key'
|
|
},
|
|
'OUTPUT_KEY': {
|
|
'value': 'img',
|
|
'description': 'input key'
|
|
},
|
|
'BACKEND': {
|
|
'value': 'pillow',
|
|
'description': 'backend'
|
|
}
|
|
}]
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
backend = cfg.get('BACKEND', BACKEND_TORCHVISION)
|
|
self.input_key = cfg.get('INPUT_KEY', 'video')
|
|
self.output_key = cfg.get('OUTPUT_KEY', 'video')
|
|
self.backend = backend
|
|
|
|
def check_video_type(self, input_video):
|
|
if self.backend == BACKEND_TORCHVISION:
|
|
assert is_tensor(input_video)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
VideoTransform.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class RandomResizedCropVideo(VideoTransform):
|
|
"""Crop a random portion of video and resize it to a given size.
|
|
|
|
Expect the video is a torch tensor with shape [..., H, W]
|
|
|
|
Args:
|
|
size (int or sequence): Desired output size.
|
|
If size is a sequence like (h, w), the output size will be matched to this.
|
|
If size is an int, the output size will be matched to (size, size).
|
|
scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
|
|
before resizing. The scale is defined with respect to the area of the original image.
|
|
ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
|
|
resizing.
|
|
interpolation (str): Desired interpolation string, 'bilinear', 'nearest', 'bicubic' are supported.
|
|
"""
|
|
para_dict = [{
|
|
'SIZE': {
|
|
'value': 0,
|
|
'description': 'size'
|
|
},
|
|
'SCALE': {
|
|
'value': [0.08, 1.0],
|
|
'description': 'scale'
|
|
},
|
|
'RATIO': {
|
|
'value': [3. / 4., 4. / 3.],
|
|
'description': 'ratio'
|
|
},
|
|
'INTERPOLATION': {
|
|
'value': 'bilinear',
|
|
'description': 'interpolation'
|
|
}
|
|
}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
size, scale = cfg.SIZE, cfg.get('SCALE', [0.08, 1.0])
|
|
ratio = cfg.get('RATIO', [3. / 4., 4. / 3.])
|
|
interpolation = cfg.get('INTERPOLATION', 'bilinear')
|
|
super(RandomResizedCropVideo, self).__init__(cfg, logger=logger)
|
|
assert self.backend in BACKENDS
|
|
self.interpolation = interpolation
|
|
if isinstance(size, (tuple, list)):
|
|
assert len(size) == 2
|
|
size = tuple(size)
|
|
elif isinstance(size, int):
|
|
size = (size, size)
|
|
else:
|
|
raise ValueError(
|
|
f'Unexpected type {type(size)}, expected int or tuple or list')
|
|
|
|
if use_video_transforms:
|
|
from torchvision.transforms._transforms_video import \
|
|
RandomResizedCropVideo as RandomResizedCropVideoOp
|
|
self.callable = RandomResizedCropVideoOp(size, scale, ratio,
|
|
self.interpolation)
|
|
else:
|
|
self.callable = transforms.RandomResizedCrop(
|
|
size, scale, ratio, INTERPOLATION_STYLE[self.interpolation])
|
|
|
|
def __call__(self, item):
|
|
self.check_video_type(item[self.input_key])
|
|
item[self.output_key] = self.callable(item[self.input_key])
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
RandomResizedCropVideo.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class CenterCropVideo(VideoTransform):
|
|
""" Crops the given video at the center.
|
|
|
|
Expect the video is a torch tensor with shape [..., H, W]
|
|
|
|
Args:
|
|
size (sequence or int): Desired output size.
|
|
If size is a sequence like (h, w), the output size will be matched to this.
|
|
If size is an int, the output size will be matched to (size, size).
|
|
"""
|
|
para_dict = [{'SIZE': {'value': 0, 'description': 'size'}}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
size = cfg.SIZE
|
|
super(CenterCropVideo, self).__init__(cfg, logger=logger)
|
|
assert self.backend in BACKENDS
|
|
self.size = size
|
|
|
|
if use_video_transforms:
|
|
from torchvision.transforms._transforms_video import \
|
|
CenterCropVideo as CenterCropVideoOp
|
|
self.callable = CenterCropVideoOp(size)
|
|
else:
|
|
self.callable = transforms.CenterCrop(size)
|
|
|
|
def __call__(self, item):
|
|
self.check_video_type(item[self.input_key])
|
|
item[self.output_key] = self.callable(item[self.input_key])
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
CenterCropVideo.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class RandomHorizontalFlipVideo(VideoTransform):
|
|
""" Horizontally flip the given video randomly with a given probability.
|
|
|
|
Expect the video is a torch tensor with shape [..., H, W]
|
|
|
|
Args:
|
|
p (float): probability of the image being flipped. Default value is 0.5
|
|
"""
|
|
para_dict = [{'P': {'value': 0.5, 'description': 'P'}}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
p = cfg.get('P', 0.5)
|
|
super(RandomHorizontalFlipVideo, self).__init__(cfg, logger=logger)
|
|
assert self.backend in BACKENDS
|
|
|
|
if use_video_transforms:
|
|
from torchvision.transforms._transforms_video import \
|
|
RandomHorizontalFlipVideo as RandomHorizontalFlipVideoOp
|
|
self.callable = RandomHorizontalFlipVideoOp(p)
|
|
else:
|
|
self.callable = transforms.RandomHorizontalFlip(p)
|
|
|
|
def __call__(self, item):
|
|
self.check_video_type(item[self.input_key])
|
|
item[self.output_key] = self.callable(item[self.input_key])
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
RandomHorizontalFlipVideo.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class NormalizeVideo(VideoTransform):
|
|
""" Normalize a tensor video with mean and standard deviation.
|
|
Expect the video is a torch tensor with shape [..., H, W]
|
|
|
|
Args:
|
|
mean (sequence): Sequence of means for each channel.
|
|
std (sequence): Sequence of standard deviations for each channel.
|
|
"""
|
|
para_dict = [{
|
|
'MEAN': {
|
|
'value': [],
|
|
'description': 'mean'
|
|
},
|
|
'STD': {
|
|
'value': [],
|
|
'description': 'std'
|
|
}
|
|
}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
mean, std = cfg.MEAN, cfg.STD
|
|
super(NormalizeVideo, self).__init__(cfg, logger=logger)
|
|
assert self.backend in BACKENDS
|
|
self.mean = np.array(mean, dtype=np.float32)
|
|
self.std = np.array(std, dtype=np.float32)
|
|
|
|
if use_video_transforms:
|
|
from torchvision.transforms._transforms_video import \
|
|
NormalizeVideo as NormalizeVideoOp
|
|
self.callable = NormalizeVideoOp(self.mean, self.std)
|
|
else:
|
|
self.callable = transforms.Normalize(self.mean, self.std)
|
|
|
|
def __call__(self, item):
|
|
video = item[self.input_key]
|
|
if not use_video_transforms:
|
|
video = video.permute(1, 0, 2, 3)
|
|
video = self.callable(video)
|
|
if not use_video_transforms:
|
|
video = video.permute(1, 0, 2, 3)
|
|
item[self.output_key] = video
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
NormalizeVideo.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class VideoToTensor(VideoTransform):
|
|
""" Convert a uint8 type tensor to a float32 tensor, permute it and scale output to [0.0, 1.0].
|
|
|
|
Expect the video is a uint8 torch tensor with shape [T, H, W, C]
|
|
"""
|
|
para_dict = [{}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(VideoToTensor, self).__init__(cfg, logger=logger)
|
|
assert self.backend in BACKENDS
|
|
|
|
def __call__(self, item):
|
|
video = item[self.input_key]
|
|
if isinstance(video, np.ndarray):
|
|
video = torch.tensor(video)
|
|
|
|
if not torch.is_tensor(video):
|
|
raise TypeError('video should be Tensor. Got %s' % type(video))
|
|
|
|
if not video.ndimension() == 4:
|
|
raise ValueError('video should be 4D. Got %dD' % video.dim())
|
|
|
|
if not video.dtype == torch.uint8:
|
|
raise TypeError(
|
|
'video tensor should have data type uint8. Got %s' %
|
|
str(video.dtype))
|
|
|
|
item[self.output_key] = video.float().permute(3, 0, 1, 2) / 255.0
|
|
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
VideoToTensor.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class AutoResizedCropVideo(VideoTransform):
|
|
""" Crop the video with a given position and resize it to a given size.
|
|
|
|
Expect the video is a torch tensor with shape [..., H, W].
|
|
|
|
Input ``crop_mode`` supports values:
|
|
- `cc`: center-center
|
|
- `cl`: left-center
|
|
- `cr`: right-center
|
|
- `tl`: left-top
|
|
- `tr`: right-top
|
|
- `bl`: left-bottom
|
|
- `br`: right-bottom
|
|
|
|
Args:
|
|
size (int or sequence): Desired output size.
|
|
If size is a sequence like (h, w), the output size will be matched to this.
|
|
If size is an int, the output size will be matched to (size, size).
|
|
scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
|
|
before resizing. The scale is defined with respect to the area of the original image.
|
|
interpolation (str): Desired interpolation string, 'bilinear', 'nearest', 'bicubic' are supported.
|
|
"""
|
|
para_dict = [{'SCALE': {'value': [0.08, 1.0], 'description': 'scale'}}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
size, scale = cfg.SIZE, cfg.get('SCALE', [0.08, 1.0])
|
|
interpolation = cfg.get('INTERPOLATION', 'bilinear')
|
|
super(AutoResizedCropVideo, self).__init__(cfg, logger=logger)
|
|
if isinstance(size, (tuple, list)):
|
|
assert len(size) == 2
|
|
size = tuple(size)
|
|
elif isinstance(size, int):
|
|
size = (size, size)
|
|
else:
|
|
raise ValueError(
|
|
f'Unexpected type {type(size)}, expected int or tuple or list')
|
|
self.size = size
|
|
self.scale = scale
|
|
self.interpolation_mode = interpolation
|
|
|
|
def get_crop(self, clip, crop_mode='cc'):
|
|
scale = random.uniform(*self.scale)
|
|
_, _, video_height, video_width = clip.shape
|
|
min_length = min(video_height, video_width)
|
|
crop_size = int(min_length * scale)
|
|
center_x = video_width // 2
|
|
center_y = video_height // 2
|
|
box_half = crop_size // 2
|
|
|
|
# default is cc
|
|
x0 = center_x - box_half
|
|
y0 = center_y - box_half
|
|
if crop_mode == 'cl':
|
|
x0 = 0
|
|
y0 = center_y - box_half
|
|
elif crop_mode == 'cr':
|
|
x0 = video_width - crop_size
|
|
y0 = center_y - box_half
|
|
elif crop_mode == 'tl':
|
|
x0 = 0
|
|
y0 = 0
|
|
elif crop_mode == 'tr':
|
|
x0 = video_width - crop_size
|
|
y0 = 0
|
|
elif crop_mode == 'bl':
|
|
x0 = 0
|
|
y0 = video_height - crop_size
|
|
elif crop_mode == 'br':
|
|
x0 = video_width - crop_size
|
|
y0 = video_height - crop_size
|
|
|
|
if use_video_transforms:
|
|
from torchvision.transforms.functional import resized_crop
|
|
return resized_crop(clip, y0, x0, crop_size, crop_size, self.size,
|
|
self.interpolation_mode)
|
|
else:
|
|
return functional.resized_crop(
|
|
clip, y0, x0, crop_size, crop_size, self.size,
|
|
INTERPOLATION_STYLE[self.interpolation_mode])
|
|
|
|
def __call__(self, item):
|
|
self.check_video_type(item[self.input_key])
|
|
crop_mode = item['meta'].get('crop_mode') or 'cc'
|
|
item[self.output_key] = self.get_crop(item[self.input_key], crop_mode)
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
AutoResizedCropVideo.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@TRANSFORMS.register_class()
|
|
class ResizeVideo(VideoTransform):
|
|
"""Resize video to a given size.
|
|
|
|
Expect the video is a torch tensor with shape [..., H, W].
|
|
|
|
Args:
|
|
size (int or sequence): Desired output size.
|
|
If size is a sequence like (h, w), the output size will be matched to this.
|
|
If size is an int, the smaller edge of the image will be matched to this
|
|
number maintaining the aspect ratio.
|
|
interpolation (str): Desired interpolation string, 'bilinear', 'nearest', 'bicubic' are supported.
|
|
"""
|
|
para_dict = [{
|
|
'SCALE': {
|
|
'value': [0.08, 1.0],
|
|
'description': 'scale'
|
|
},
|
|
'INTERPOLATION': {
|
|
'value': 'bilinear',
|
|
'description': 'interpolation'
|
|
}
|
|
}]
|
|
para_dict[0].update(VideoTransform.para_dict[0])
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
size = cfg.SIZE
|
|
interpolation = cfg.get('INTERPOLATION', 'bilinear')
|
|
super(ResizeVideo, self).__init__(cfg, logger=logger)
|
|
self.size = size
|
|
if isinstance(self.size, (tuple, list)):
|
|
self.size = tuple(self.size)
|
|
assert len(self.size) == 2
|
|
else:
|
|
if not isinstance(self.size, int):
|
|
raise ValueError(
|
|
f'Expected size to be tuple or list or int, got {type(self.size)}'
|
|
)
|
|
self.interpolation_mode = interpolation
|
|
|
|
def resize(self, clip):
|
|
if use_video_transforms:
|
|
from torchvision.transforms.functional import resize
|
|
|
|
# resize function only takes a tuple size
|
|
# so we need to compute scaled target size here
|
|
if isinstance(self.size, int):
|
|
h, w = clip.shape[-2], clip.shape[-1]
|
|
if (w <= h and w == self.size) or (h <= w and h == self.size):
|
|
return clip
|
|
if w < h:
|
|
ow = self.size
|
|
oh = int(self.size * h / w)
|
|
else:
|
|
oh = self.size
|
|
ow = int(self.size * w / h)
|
|
size = (oh, ow)
|
|
else:
|
|
size = self.size
|
|
return resize(clip, size, self.interpolation_mode)
|
|
else:
|
|
return functional.resize(
|
|
clip, self.size, INTERPOLATION_STYLE[self.interpolation_mode])
|
|
|
|
def __call__(self, item):
|
|
self.check_video_type(item[self.input_key])
|
|
item[self.output_key] = self.resize(item[self.input_key])
|
|
return item
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('TRANSFORM',
|
|
__class__.__name__,
|
|
ResizeVideo.para_dict,
|
|
set_name=True)
|