| 115 | return self._rand_range(*self.factor_range) |
| 116 | |
| 117 | def _augment(self, img, r): |
| 118 | old_dtype = img.dtype |
| 119 | |
| 120 | if img.ndim == 3: |
| 121 | if self.rgb is not None: |
| 122 | m = cv2.COLOR_RGB2GRAY if self.rgb else cv2.COLOR_BGR2GRAY |
| 123 | grey = cv2.cvtColor(img.astype('float32'), m) |
| 124 | mean = np.mean(grey) |
| 125 | else: |
| 126 | mean = np.mean(img, axis=(0, 1), keepdims=True) |
| 127 | else: |
| 128 | mean = np.mean(img) |
| 129 | |
| 130 | img = img * r + mean * (1 - r) |
| 131 | if self.clip or old_dtype == np.uint8: |
| 132 | img = np.clip(img, 0, 255) |
| 133 | return img.astype(old_dtype) |
| 134 | |
| 135 | |
| 136 | class MeanVarianceNormalize(PhotometricAugmentor): |