Merge pull request #38 from arsenyinfo/jpeg_with_float32

supporting float32 input for jpeg augmentation

albumentations/augmentations/functional.py

@@ -1,8 +1,9 @@
 from __future__ import division
+
 from functools import wraps
+from warnings import warn
 
 import cv2
-
 import numpy as np
 from scipy.ndimage.filters import gaussian_filter
 
@@ -214,10 +215,24 @@ def motion_blur(img, ksize):
 
 
 def jpeg_compression(img, quality):
-    if img.dtype != np.uint8:
-        raise TypeError('jpeg_compression supports only uint8 inputs')
+    input_dtype = img.dtype
+    needs_float = False
+
+    if input_dtype == np.float32:
+        warn('Jpeg compression augmentation '
+             'is most effective with uint8 inputs, '
+             '{} is used as input.'.format(input_dtype),
+             UserWarning)
+        img = from_float(img, dtype=np.dtype('uint8'))
+        needs_float = True
+    elif input_dtype not in (np.uint8, np.float32):
+        raise ValueError('Unexpected dtype {} for Jpeg augmentation'.format(input_dtype))
+
     _, encoded_img = cv2.imencode('.jpg', img, (cv2.IMWRITE_JPEG_QUALITY, quality))
     img = cv2.imdecode(encoded_img, cv2.IMREAD_UNCHANGED)
+
+    if needs_float:
+        img = to_float(img, max_value=255)
     return img
 
 

albumentations/augmentations/transforms.py

@@ -244,6 +244,11 @@ class CenterCrop(DualTransform):
 
     Image types:
         uint8, float32
+
+    Note:
+        It is recommended to use uint8 images as input.
+        Otherwise the operation will require internal conversion
+        float32 -> uint8 -> float32 that causes worse performance.
     """
 
     def __init__(self, height, width, p=1.0):
@@ -431,7 +436,7 @@ class JpegCompression(ImageOnlyTransform):
             image
 
         Image types:
-            uint8
+            uint8, float32
         """
 
     def __init__(self, quality_lower=99, quality_upper=100, p=0.5):

setup.py

@@ -34,7 +34,7 @@ def get_long_description():
     license='MIT',
     url='https://github.com/albu/albumentations',
     packages=find_packages(exclude=['tests']),
-    install_requires=['numpy>=1.11.1', 'scipy', 'opencv-python', 'imgaug'],
+    install_requires=['numpy>=1.11.1', 'scipy', 'opencv-python', 'imgaug>=0.2.5'],
     extras_require={'tests': get_test_requirements()},
     classifiers=[
         'Development Status :: 4 - Beta',

tests/test_augmentations.py

@@ -53,6 +53,7 @@ def test_image_only_augmentations(augmentation_cls, params, image, mask):
     [ChannelShuffle, {}],
     [InvertImg, {}],
     [RandomGamma, {}],
+    [JpegCompression, {}],
     [ToGray, {}],
     [Cutout, {}],
 ])