Unverified Commit 7dac84d0 authored by Rui Xu's avatar Rui Xu Committed by GitHub
Browse files

support pathlib in imageio (#208)



* support pathlib in imageio

* support pathlib in imageio

* add Path type for the annotation.
Co-authored-by: default avatarxurui4 <xurui4@sensetime.com>
parent c294d271
# Copyright (c) Open-MMLab. All rights reserved.
import os.path as osp
from pathlib import Path
import cv2
import numpy as np
......@@ -61,9 +62,9 @@ def imread(img_or_path, flag='color', channel_order='bgr'):
"""Read an image.
Args:
img_or_path (ndarray or str): Either a numpy array or image path.
If it is a numpy array (loaded image), then it will be returned
as is.
img_or_path (ndarray or str or Path): Either a numpy array or str or
pathlib.Path. If it is a numpy array (loaded image), then
it will be returned as is.
flag (str): Flags specifying the color type of a loaded image,
candidates are `color`, `grayscale` and `unchanged`.
Note that the `turbojpeg` backened does not support `unchanged`.
......@@ -72,6 +73,9 @@ def imread(img_or_path, flag='color', channel_order='bgr'):
Returns:
ndarray: Loaded image array.
"""
if isinstance(img_or_path, Path):
img_or_path = str(img_or_path)
if isinstance(img_or_path, np.ndarray):
return img_or_path
elif is_str(img_or_path):
......@@ -91,7 +95,8 @@ def imread(img_or_path, flag='color', channel_order='bgr'):
cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
return img
else:
raise TypeError('"img" must be a numpy array or a filename')
raise TypeError('"img" must be a numpy array or a str or '
'a pathlib.Path object')
def imfrombytes(content, flag='color', channel_order='bgr'):
......
......@@ -2,6 +2,7 @@
import os
import os.path as osp
import tempfile
from pathlib import Path
import cv2
import numpy as np
......@@ -17,8 +18,10 @@ class TestImage(object):
def setup_class(cls):
# the test img resolution is 400x300
cls.img_path = osp.join(osp.dirname(__file__), 'data/color.jpg')
cls.img_path_obj = Path(cls.img_path)
cls.gray_img_path = osp.join(
osp.dirname(__file__), 'data/grayscale.jpg')
cls.gray_img_path_obj = Path(cls.gray_img_path)
cls.img = cv2.imread(cls.img_path)
cls.mean = np.float32(np.array([123.675, 116.28, 103.53]))
cls.std = np.float32(np.array([58.395, 57.12, 57.375]))
......@@ -47,6 +50,18 @@ class TestImage(object):
assert img_cv2_unchanged.shape == (300, 400)
img_cv2_unchanged = mmcv.imread(img_cv2_unchanged)
assert_array_equal(img_cv2_unchanged, mmcv.imread(img_cv2_unchanged))
img_cv2_color_bgr = mmcv.imread(self.img_path_obj)
assert img_cv2_color_bgr.shape == (300, 400, 3)
img_cv2_color_rgb = mmcv.imread(self.img_path_obj, channel_order='rgb')
assert img_cv2_color_rgb.shape == (300, 400, 3)
assert_array_equal(img_cv2_color_rgb[:, :, ::-1], img_cv2_color_bgr)
img_cv2_grayscale1 = mmcv.imread(self.img_path_obj, 'grayscale')
assert img_cv2_grayscale1.shape == (300, 400)
img_cv2_grayscale2 = mmcv.imread(self.gray_img_path_obj)
assert img_cv2_grayscale2.shape == (300, 400, 3)
img_cv2_unchanged = mmcv.imread(self.gray_img_path_obj, 'unchanged')
assert img_cv2_unchanged.shape == (300, 400)
with pytest.raises(TypeError):
mmcv.imread(1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment