File size: 899 Bytes
66655aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
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 |