preprocesses.py

import cv2
import numpy as np
import plasma.functional as f
import albumentations as alb


class Preprocessor(f.SequentialPipe):

    def __init__(self, width, interpolation):
        super().__init__()

        self.resize = WidthReisze(width, interpolation)
        # self.to_gray = ToGray()


class WidthReisze(f.Pipe):

    def __init__(self, width, interpolation):
        super().__init__(width=width, interpolation=interpolation)
    
    def run(self, img):
        w = self.width
        resized_img = cv2.resize(img, (w, int(w / img.shape[1] * img.shape[0])), interpolation=self.interpolation)
        return resized_img


class ToGray(f.Pipe):

    def __init__(self):
        super().__init__()

    def run(self, img):
        if len(img.shape) == 3:
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        return img