瀏覽代碼

Image reading is now done with plain PIL. This enables lazy image reading from disc

Dimitri Korsch 6 年之前
父節點
當前提交
a568ef7d90
共有 4 個文件被更改,包括 75 次插入19 次删除
  1. 1 1
      nabirds/__init__.py
  2. 46 9
      nabirds/dataset/image.py
  3. 27 8
      nabirds/dataset/utils.py
  4. 1 1
      requirements.txt

+ 1 - 1
nabirds/__init__.py

@@ -1,4 +1,4 @@
 from .dataset import Dataset
 from .annotations import NAB_Annotations, CUB_Annotations
 
-__version__ = "0.1.8"
+__version__ = "0.2.0"

+ 46 - 9
nabirds/dataset/image.py

@@ -1,4 +1,5 @@
 from imageio import imread
+from PIL import Image
 from os.path import isfile
 
 import copy
@@ -13,12 +14,18 @@ def should_have_parts(func):
 	return inner
 
 class ImageWrapper(object):
+	@staticmethod
+	def read_image(im_path, mode="RGB"):
+		# im = imread(im_path, pilmode=mode)
+		im = Image.open(im_path, mode="r")
+		return im
+
+
 	def __init__(self, im_path, label, parts=None, mode="RGB"):
-		if isinstance(im_path, str):
-			assert isfile(im_path), "Image \"{}\" does not exist!".format(im_path)
-			self.im = imread(im_path, pilmode=mode)
-		else:
-			self.im = im_path
+
+		self.mode = mode
+		self.im = im_path
+		self._im_array = None
 
 		self.label = label
 		self.parts = parts
@@ -26,14 +33,43 @@ class ImageWrapper(object):
 		self.parent = None
 		self._feature = None
 
+	def __del__(self):
+		self._im.close()
+
+	@property
+	def im(self):
+		return self._im
+
+	@property
+	def im_array(self):
+		if self._im_array is None:
+			self.im = self.im.convert(self.mode)
+			self._im_array = utils.asarray(self.im)
+		return self._im_array
+
+	@im.setter
+	def im(self, value):
+		if isinstance(value, str):
+			assert isfile(value), "Image \"{}\" does not exist!".format(value)
+			self._im = ImageWrapper.read_image(value, mode=self.mode)
+		else:
+			self._im = value
+
 	def as_tuple(self):
-		return self.im, self.parts, self.label
+		return self.im_array, self.parts, self.label
 
 	def copy(self):
-		new = copy.deepcopy(self)
+		new = copy.copy(self)
 		new.parent = self
-		return new
+		deepcopies = [
+			"_feature",
+			"parts",
+		]
+		for attr_name in deepcopies:
+			attr_copy = copy.deepcopy(getattr(self, attr_name))
+			setattr(new, attr_name, attr_copy)
 
+		return new
 	@property
 	def feature(self):
 		return self._feature
@@ -44,7 +80,8 @@ class ImageWrapper(object):
 
 	def crop(self, x, y, w, h):
 		result = self.copy()
-		result.im = self.im[y:y+h, x:x+w]
+		# result.im = self.im[y:y+h, x:x+w]
+		result.im = self.im.crop(x, y, x+w, y+h)
 		if self.has_parts:
 			result.parts[:, 1] -= x
 			result.parts[:, 2] -= y

+ 27 - 8
nabirds/dataset/utils.py

@@ -1,13 +1,34 @@
 import numpy as np
+from PIL.Image import Image as PIL_Image
 
 DEFAULT_RATIO = np.sqrt(49 / 400)
 
 def __expand_parts(p):
 	return p[:, 0], p[:, 1:3], p[:, 3].astype(bool)
 
+def dimensions(im):
+	if isinstance(im, np.ndarray):
+		assert im.ndim == 3, "Only RGB images are currently supported!"
+		return im.shape
+	elif isinstance(im, PIL_Image):
+		w, h = im.size
+		c = len(im.getbands())
+		# assert c == 3, "Only RGB images are currently supported!"
+		return h, w, c
+	else:
+		raise ValueError("Unknown image instance ({})!".format(type(im)))
+
+def asarray(im, dtype=np.uint8):
+	if isinstance(im, np.ndarray):
+		return im.astype(dtype)
+	elif isinstance(im, PIL_Image):
+		return np.asarray(im, dtype=dtype)
+	else:
+		raise ValueError("Unknown image instance ({})!".format(type(im)))
+
 
 def uniform_parts(im, ratio=DEFAULT_RATIO, round_op=np.floor):
-	h, w, c = im.shape
+	h, w, c = dimensions(im)
 
 	part_w = round_op(w * ratio).astype(np.int32)
 	part_h = round_op(h * ratio).astype(np.int32)
@@ -31,11 +52,9 @@ def visible_part_locs(p):
 
 
 def crops(im, xy, ratio=DEFAULT_RATIO, padding_mode="edge"):
-	assert im.ndim == 3, "Only RGB images are currently supported!"
-
-	h, w, c = im.shape
+	h, w, c = dimensions(im)
 	crop_h, crop_w = int(h * ratio), int(w * ratio)
-	crops = np.zeros((xy.shape[1], crop_h, crop_w, c), dtype=im.dtype)
+	crops = np.zeros((xy.shape[1], crop_h, crop_w, c), dtype=np.uint8)
 
 	pad_h, pad_w = crop_h // 2, crop_w // 2
 
@@ -49,15 +68,15 @@ def crops(im, xy, ratio=DEFAULT_RATIO, padding_mode="edge"):
 
 def visible_crops(im, p, *args, **kw):
 	idxs, locs, vis = __expand_parts(p)
-	parts = crops(im, locs[vis].T, *args, **kw)
+	parts = crops(asarray(im), locs[vis].T, *args, **kw)
 	res = np.zeros((len(idxs),) + parts.shape[1:], dtype=parts.dtype)
 	res[vis] = parts
 	return res
 
 def reveal_parts(im, xy, ratio=DEFAULT_RATIO):
-	h, w, c = im.shape
+	h, w, c = dimensions(im)
 	crop_h, crop_w = int(h * ratio), int(w * ratio)
-
+	im = asarray(im)
 	res = np.zeros_like(im)
 	for x, y in xy.T:
 		x0, y0 = max(x - crop_w // 2, 0), max(y - crop_h // 2, 0)

+ 1 - 1
requirements.txt

@@ -1,4 +1,4 @@
-imageio
+#imageio
 numpy
 pillow
 matplotlib