Diff of /data/data_utils.py [000000] .. [bd7f9c]

Switch to side-by-side view

--- a
+++ b/data/data_utils.py
@@ -0,0 +1,86 @@
+import random
+import numpy as np
+from PIL import Image
+
+def read_image(path, dtype=np.float32, color=True):
+    """
+    Read an image from a given path, the image is in (C, H, W) format and the range of its value is between [0, 255]
+    :param path: The path of an image file
+    :param dtype: data type of an image, default is float32
+    :param color: If 'True', the number of channels is three, in this case, RGB
+        If 'False', this function returns a gray scale image
+    :return:
+    """
+    f = Image.open(path)
+    try:
+        if color:
+            img = f.convert('RGB')
+        else:
+            img = f.convert('P')
+        img = np.asarray(img, dtype=dtype)
+    finally:
+        if hasattr(f, 'close'):
+            f.close()
+
+    if img.ndim == 2:
+        # reshape (H, W) -> (1, H, W)
+        return img[np.newaxis]
+    else:
+        # transpose (H, W, C) -> (C, H, W)
+        return img.transpose((2, 0, 1))
+
+
+def flip_masks(masks, y_flip=False, x_flip=False):
+    masks = masks.copy()
+    for i in range(masks.shape[0]):
+        if y_flip:
+            masks[i] = np.flipud(masks[i]).copy()
+        if x_flip:
+            masks[i] = np.fliplr(masks[i]).copy()
+    return masks
+
+def random_flip(img, y_random=False, x_random=False,
+                return_param=False, copy=False):
+    """
+    Randomly flip an image in vertical or horizontal direction.
+    :param img (numpy.ndarray): An array that gets flipped.
+        This is in CHW format.
+    :param y_random (bool): Randomly flip in vertical direction.
+    :param x_random (bool): Randomly flip in horizontal direction.
+    :param return_param (bool): Returns information of flip.
+    :param copy (bool): If False, a view of :obj:`img` will be returned.
+    :return (numpy.ndarray or (numpy.ndarray, dict)):
+        If :`return_param = False`,
+        returns an array :obj:`out_img` that is the result of flipping.
+        If :obj:`return_param = True`,
+        returns a tuple whose elements are :obj:`out_img, param`.
+        :obj:`param` is a dictionary of intermediate parameters whose
+        contents are listed below with key, value-type and the description
+        of the value.
+        * **y_flip** (*bool*): Whether the image was flipped in the\
+            vertical direction or not.
+        * **x_flip** (*bool*): Whether the image was flipped in the\
+            horizontal direction or not.
+    """
+    y_flip, x_flip = False, False
+    if y_random:
+        y_flip = random.choice([True, False])
+    if x_random:
+        x_flip = random.choice([True, False])
+
+    if y_flip:
+        img = img[:, ::-1, :]
+    if x_flip:
+        img = img[:, :, ::-1]
+
+    if copy:
+        img = img.copy()
+
+    if return_param:
+        return img, {'y_flip': y_flip, 'x_flip': x_flip}
+    else:
+        return img
+
+
+if __name__ == '__main__':
+    pass
\ No newline at end of file