Preprocessing under cv2 conventions.
| 19 | |
| 20 | |
| 21 | class Preprocessor: |
| 22 | |
| 23 | """ |
| 24 | Preprocessing under cv2 conventions. |
| 25 | """ |
| 26 | |
| 27 | def __init__(self): |
| 28 | self.rembg_session = rembg.new_session( |
| 29 | providers=["CUDAExecutionProvider", "CPUExecutionProvider"], |
| 30 | ) |
| 31 | |
| 32 | def preprocess(self, image_path: str, save_path: str, rmbg: bool = True, recenter: bool = True, size: int = 512, border_ratio: float = 0.2): |
| 33 | image = self.step_load_to_size(image_path=image_path, size=size*2) |
| 34 | if rmbg: |
| 35 | image = self.step_rembg(image_in=image) |
| 36 | else: |
| 37 | image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA) |
| 38 | if recenter: |
| 39 | image = self.step_recenter(image_in=image, border_ratio=border_ratio, square_size=size) |
| 40 | else: |
| 41 | image = cv2.resize( |
| 42 | src=image, |
| 43 | dsize=(size, size), |
| 44 | interpolation=cv2.INTER_AREA, |
| 45 | ) |
| 46 | return cv2.imwrite(save_path, image) |
| 47 | |
| 48 | def step_rembg(self, image_in: np.ndarray) -> np.ndarray: |
| 49 | image_out = rembg.remove( |
| 50 | data=image_in, |
| 51 | session=self.rembg_session, |
| 52 | ) |
| 53 | return image_out |
| 54 | |
| 55 | def step_recenter(self, image_in: np.ndarray, border_ratio: float, square_size: int) -> np.ndarray: |
| 56 | assert image_in.shape[-1] == 4, "Image to recenter must be RGBA" |
| 57 | mask = image_in[..., -1] > 0 |
| 58 | ijs = np.nonzero(mask) |
| 59 | # find bbox |
| 60 | i_min, i_max = ijs[0].min(), ijs[0].max() |
| 61 | j_min, j_max = ijs[1].min(), ijs[1].max() |
| 62 | bbox_height, bbox_width = i_max - i_min, j_max - j_min |
| 63 | # recenter and resize |
| 64 | desired_size = int(square_size * (1 - border_ratio)) |
| 65 | scale = desired_size / max(bbox_height, bbox_width) |
| 66 | desired_height, desired_width = int(bbox_height * scale), int(bbox_width * scale) |
| 67 | desired_i_min, desired_j_min = (square_size - desired_height) // 2, (square_size - desired_width) // 2 |
| 68 | desired_i_max, desired_j_max = desired_i_min + desired_height, desired_j_min + desired_width |
| 69 | # create new image |
| 70 | image_out = np.zeros((square_size, square_size, 4), dtype=np.uint8) |
| 71 | image_out[desired_i_min:desired_i_max, desired_j_min:desired_j_max] = cv2.resize( |
| 72 | src=image_in[i_min:i_max, j_min:j_max], |
| 73 | dsize=(desired_width, desired_height), |
| 74 | interpolation=cv2.INTER_AREA, |
| 75 | ) |
| 76 | return image_out |
| 77 | |
| 78 | def step_load_to_size(self, image_path: str, size: int) -> np.ndarray: |
no outgoing calls
no test coverage detected