Extract alpha and resize to original. Returns alpha uint8.
(self, pred: np.ndarray, original_size: tuple)
| 145 | return img_batch, (h, w) |
| 146 | |
| 147 | def _postprocess(self, pred: np.ndarray, original_size: tuple) -> np.ndarray: |
| 148 | """Extract alpha and resize to original. Returns alpha uint8.""" |
| 149 | # Remove batch, get alpha |
| 150 | alpha = pred[0, 0, :, :] |
| 151 | |
| 152 | # Resize to original |
| 153 | alpha_resized = cv2.resize(alpha, (original_size[1], original_size[0]), interpolation=cv2.INTER_LINEAR) |
| 154 | |
| 155 | # To uint8 |
| 156 | alpha_resized = (alpha_resized * 255).astype(np.uint8) |
| 157 | |
| 158 | return alpha_resized |
| 159 | |
| 160 | def predict(self, image: Image.Image) -> Image.Image: |
| 161 | """Background removal; fallback to CPU if GPU fails. Returns RGBA PIL.""" |