81 lines
3.2 KiB
Python
81 lines
3.2 KiB
Python
from typing import Any, Dict, Tuple
|
|
import torchvision.transforms as transforms
|
|
|
|
|
|
class PtvTransformsDataAugment:
|
|
"""
|
|
Applies data augmentation transformations to dataset elements.
|
|
Supported transformations include flipping, rotation, translation, and scaling.
|
|
|
|
category: PyTorch wrapper - Training
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the required input types for the transformations.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary specifying the required input types and their constraints.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"h_flip_prob": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"v_flip_prob": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"rotate_degree": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.001}),
|
|
"h_translate_ratio": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"v_translate_ratio": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"min_scale": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 1.0, "step": 0.001}),
|
|
"max_scale": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.001})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str] = ("PTVTRANSFORM",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Training"
|
|
|
|
def f(self,
|
|
h_flip_prob: float,
|
|
v_flip_prob: float,
|
|
rotate_degree: float,
|
|
h_translate_ratio: float,
|
|
v_translate_ratio: float,
|
|
min_scale: float,
|
|
max_scale: float) -> Tuple[transforms.Compose]:
|
|
"""
|
|
Applies the specified transformations to the dataset.
|
|
|
|
Args:
|
|
h_flip_prob (float): Probability of applying horizontal flip.
|
|
v_flip_prob (float): Probability of applying vertical flip.
|
|
rotate_degree (float): Maximum rotation degree.
|
|
h_translate_ratio (float): Horizontal translation ratio.
|
|
v_translate_ratio (float): Vertical translation ratio.
|
|
min_scale (float): Minimum scale factor.
|
|
max_scale (float): Maximum scale factor.
|
|
|
|
Returns:
|
|
Tuple[transforms.Compose]: A tuple containing the transformation pipeline.
|
|
"""
|
|
translate_tuple = (h_translate_ratio, v_translate_ratio) if h_translate_ratio > 0.0 or v_translate_ratio > 0.0 else None
|
|
scale_tuple = (min_scale, max_scale) if min_scale != 1.0 or max_scale != 1.0 else None
|
|
|
|
actions = []
|
|
if h_flip_prob > 0.0:
|
|
actions.append(transforms.RandomHorizontalFlip(p=h_flip_prob))
|
|
if v_flip_prob > 0.0:
|
|
actions.append(transforms.RandomVerticalFlip(p=v_flip_prob))
|
|
|
|
actions.append(
|
|
transforms.RandomAffine(
|
|
degrees=rotate_degree,
|
|
translate=translate_tuple,
|
|
scale=scale_tuple
|
|
)
|
|
)
|
|
actions.append(transforms.ToTensor()) # Converts image to tensor and scales to [0, 1]
|
|
|
|
transform_pipeline = transforms.Compose(actions)
|
|
|
|
return (transform_pipeline,)
|