Pre-Processing

├── custom
│   ├── __init__.py
│   ├── custom_gray2rgb.py
├── train.yaml
└── transforms.yaml

from avi_transforms import DataItem

class Gray2RGB:
  """Test transform to convert Gray to RGB"""

    def __init__(self, defect_map, **kwargs):
        # Example to show passing arguments to custom transform
        print(defect_map)

    def __call__(self, inputs: DataItem) -> DataItem:
        image = inputs.image
        if image.ndim == 3:
            image = image[..., 0]
        rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        return DataItem(image=rgb)

train:
  - GaussianBlur:
      lower_limit: 5
      upper_limit: 9
      p: 0.8
  - VerticalFlip:
      p: 0.2
  - CustomTransform:
      transform: custom.custom_gray2rgb.Gray2RGB
      params:
        defect_map:
          0: ok
          1: ng
      p: 1.0
valid: