fix when opencv reshapes 1-channel image (fix #184)

tensorpack/dataflow/imgaug/geometry.py

@@ -37,6 +37,8 @@ class Shift(ImageAugmentor):
     def _augment(self, img, shift_m):
         ret = cv2.warpAffine(img, shift_m, img.shape[1::-1],
                              borderMode=self.border)
+        if img.ndim == 3 and ret.ndim == 2:
+            ret = ret[:, :, np.newaxis]
         return ret
 
 
@@ -71,6 +73,8 @@ class Rotation(ImageAugmentor):
     def _augment(self, img, rot_m):
         ret = cv2.warpAffine(img, rot_m, img.shape[1::-1],
                              flags=self.interp, borderMode=self.border)
+        if img.ndim == 3 and ret.ndim == 2:
+            ret = ret[:, :, np.newaxis]
         return ret
 
 
@@ -99,6 +103,8 @@ class RotationAndCropValid(ImageAugmentor):
         rot_m = cv2.getRotationMatrix2D((center[0] - 0.5, center[1] - 0.5), deg, 1)
         ret = cv2.warpAffine(img, rot_m, img.shape[1::-1],
                              flags=self.interp, borderMode=cv2.BORDER_CONSTANT)
+        if img.ndim == 3 and ret.ndim == 2:
+            ret = ret[:, :, np.newaxis]
         neww, newh = RotationAndCropValid.largest_rotated_rect(ret.shape[1], ret.shape[0], deg)
         neww = min(neww, ret.shape[1])
         newh = min(newh, ret.shape[0])

tensorpack/dataflow/imgaug/imgproc.py

@@ -152,8 +152,10 @@ class Gamma(ImageAugmentor):
         old_dtype = img.dtype
         lut = ((np.arange(256, dtype='float32') / 255) ** (1. / (1. + gamma)) * 255).astype('uint8')
         img = np.clip(img, 0, 255).astype('uint8')
-        img = cv2.LUT(img, lut).astype(old_dtype)
-        return img
+        ret = cv2.LUT(img, lut).astype(old_dtype)
+        if img.ndim == 3 and ret.ndim == 2:
+            ret = ret[:, :, np.newaxis]
+        return ret
 
 
 class Clip(ImageAugmentor):

tensorpack/dataflow/imgaug/noname.py

@@ -39,8 +39,12 @@ class Flip(ImageAugmentor):
 
     def _augment(self, img, do):
         if do:
-            img = cv2.flip(img, self.code)
-        return img
+            ret = cv2.flip(img, self.code)
+            if img.ndim == 3 and ret.ndim == 2:
+                ret = ret[:, :, np.newaxis]
+        else:
+            ret = img
+        return ret
 
     def _fprop_coord(self, coord, param):
         raise NotImplementedError()