Offline augmentation

[MinhxNguyen7]

Description

• Implement offline augmentation to plug into the pipeline.
• The augmentations will be probabilistically applied.
• Research augmentation techniques and see which one would be useful.

Starting Code

from abc import ABC, abstractmethod
from copy import deepcopy
from pathlib import Path
import random
from typing import Iterable
from yolo_to_yolo.data_transformers import YoloDataTransformer
from yolo_to_yolo.data_types import YoloImageData
from yolo_to_yolo.yolo_data_pipeline import YoloDataPipeline


class Augmentation(ABC):
    @abstractmethod
    def __call__(self, input_data: YoloImageData) -> YoloImageData:
        ...


class AugmentationPipeline(YoloDataTransformer):
    def __init__(self, augmentations: list[Augmentation], probabilities: list[float]):
        self.augmentations = augmentations
        self.probabilities = probabilities
        
    def __call__(self, input_data: YoloImageData) -> Iterable[YoloImageData]:
        for _ in range(10):
            yield self._apply_augmentations(deepcopy(input_data))      
            
    def _apply_augmentations(self, input_data: YoloImageData) -> YoloImageData:
        data = input_data
        for augmentation, probability in zip(self.augmentations, self.probabilities):
            if random.random() < probability:
                data = augmentation(data)
        
        return data
    

class RandomFlip(Augmentation):
    def __call__(self, input_data: YoloImageData) -> YoloImageData:
        # TODO: Implement
        ...
        

class DirectionalBlur(Augmentation):
    def __call__(self, input_data: YoloImageData) -> YoloImageData:
        # TODO: Implement
        ...


if __name__ == '__main__':
    """
    Example usage
    """
    augmentations = AugmentationPipeline(
        [RandomFlip(), DirectionalBlur()],
        [0.5, 0.3]
    )
    
    pipeline = YoloDataPipeline([augmentations])
    
    pipeline.apply_to_dir(
        input_dir=Path('input_dir'),
        output_dir=Path('output_dir')
    )