Fix contrast mean value for float32 and factor=0 (#1207)

albumentations-team · Jul 9, 2022 · ca59440 · ca59440
1 parent a4d33e1
commit ca59440
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py
@@ -1771,7 +1771,9 @@ def adjust_contrast_torchvision(img, factor):
         mean = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).mean()
 
     if factor == 0:
-        return np.full_like(img, int(mean + 0.5), dtype=img.dtype)
+        if img.dtype != np.float32:
+            mean = int(mean + 0.5)
+        return np.full_like(img, mean, dtype=img.dtype)
 
     if img.dtype == np.uint8:
         return _adjust_contrast_torchvision_uint8(img, factor, mean)