created new class for part annotations handling

nabirds/dataset/image.py

@@ -6,6 +6,7 @@ import copy
 import numpy as np
 
 from . import utils
+from .part import Parts
 
 def should_have_parts(func):
 	def inner(self, *args, **kwargs):
@@ -23,15 +24,12 @@ class ImageWrapper(object):
 
 	def __init__(self, im_path, label, parts=None, mode="RGB", part_rescale_size=None):
 
-
 		self.mode = mode
 		self.im = im_path
 		self._im_array = None
 
 		self.label = label
-		self.parts = utils.rescale_parts(self.im, parts, part_rescale_size)
-
-		self.part_rescale_size = part_rescale_size
+		self.parts = Parts(self.im, parts, part_rescale_size)
 
 		self.parent = None
 		self._feature = None
@@ -110,39 +108,37 @@ class ImageWrapper(object):
 		idxs, (xs,ys) = self.visible_part_locs()
 		f = np.logical_and
 		mask = f(f(x <= xs, xs <= x+w), f(y <= ys, ys <= y+h))
+		idxs = np.where(mask)
+		import pdb; pdb.set_trace()
 		result = self.copy()
-		result.parts[:, -1] = mask.astype(self.parts.dtype)
-
+		result.parts.select(idxs)
 		return result
 
 	def uniform_parts(self, ratio):
 		result = self.copy()
+		raise NotImplementedError("FIX ME!")
 		result.parts = utils.uniform_parts(self.im, ratio=ratio)
 		return result
 
 	@should_have_parts
 	def select_parts(self, idxs):
 		result = self.copy()
-
-		result.parts[:, -1] = 0
-		result.parts[idxs, -1] = 1
-
+		result.parts.select(idxs)
 		return result
 
 	@should_have_parts
 	def select_random_parts(self, rnd, n_parts):
-
 		idxs, xy = self.visible_part_locs()
 		rnd_idxs = utils.random_idxs(idxs, rnd=rnd, n_parts=n_parts)
 		return self.select_parts(rnd_idxs)
 
 	@should_have_parts
 	def visible_crops(self, ratio):
-		return utils.visible_crops(self.im, self.parts, ratio=ratio)
+		return self.parts.visible_crops(ratio=ratio)
 
 	@should_have_parts
 	def visible_part_locs(self):
-		return utils.visible_part_locs(self.parts)
+		return self.parts.visible_locs()
 
 	@should_have_parts
 	def reveal_visible(self, ratio):

nabirds/dataset/part.py

@@ -0,0 +1,90 @@
+import numpy as np
+
+from . import utils
+
+class Parts(object):
+	def __init__(self, image, part_annotations, rescale_size):
+		super(Parts, self).__init__()
+		annots = utils.rescale_parts(image, part_annotations, rescale_size)
+
+		self._parts = [ImagePart(image, a) for a in annots]
+		self.rescale_size = rescale_size
+
+	def __getitem__(self, i):
+		return self._parts[i]
+
+	@property
+	def selected(self):
+		return np.array([p.is_visible for p in self._parts], dtype=bool)
+
+	@property
+	def selected_idxs(self):
+		return np.where(self.selected)[0]
+
+	def select(self, idxs):
+		if isinstance(idxs, np.ndarray) and idxs.dtype == bool:
+			# a mask is present, so convert it to indeces
+			idxs = np.where(idxs)[0]
+
+		for p in self._parts:
+			p.is_visible = p._id in idxs
+
+	def invert_selection(self):
+		self.select(np.logical_not(self.selected))
+
+	def set_visibility(self, idxs, value):
+		for p in self._parts[idxs]:
+			p.is_visible = value
+
+	def visible_locs(self):
+		vis = [(p._id, p.xy) for p in self._parts if p.is_visible]
+		idxs, xy = zip(*vis)
+		return np.array(idxs), np.array(xy).T
+
+	def visible_crops(self, *args, **kwargs):
+		return np.array([p.crop(*args, **kwargs) for p in self._parts])
+
+class ImagePart(object):
+	def __init__(self, image, annotation):
+		super(ImagePart, self).__init__()
+		self.image = image
+
+		if len(annotation) == 4:
+			# here x,y are the center of the part
+			self._id, self.x, self.y, self._is_visible = annotation
+			self.w, self.h = None, None
+
+		elif len(annotation) == 5:
+			# here x,y are top left corner of the part
+			self._id, self.x, self.y, self.w, self.h = annotation
+			self._is_visible = True
+
+		else:
+			raise ValueError("Unknown annotation format: {}".format(annotation))
+
+	def crop(self, ratio=None, padding_mode="edge"):
+		if not self.is_visible:
+			h, w, c = utils.dimensions(self.image)
+			crop_h, crop_w = int(h * ratio), int(w * ratio)
+			return np.zeros((crop_h, crop_w, c), dtype=np.uint8)
+		else:
+			return utils.crop(self.image, self.xy, ratio, padding_mode)
+
+	@property
+	def is_visible(self):
+		return bool(self._is_visible)
+
+	@is_visible.setter
+	def is_visible(self, value):
+		self._is_visible = bool(value)
+
+	@property
+	def xy(self):
+		return np.array([self.x, self.y])
+
+	@property
+	def wh(self):
+		return np.array([self.w, self.h])
+
+
+

nabirds/dataset/utils.py

@@ -3,8 +3,8 @@ 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 __expand_parts(p):
+# 	return p[:, 0], p[:, 1:3], p[:, 3].astype(bool)
 
 def rescale_parts(im, parts, part_rescale_size):
 	if part_rescale_size is None or part_rescale_size < 0:
@@ -66,32 +66,37 @@ def uniform_parts(im, ratio=DEFAULT_RATIO, round_op=np.floor):
 	return parts
 
 def visible_part_locs(p):
-	idxs, locs, vis = __expand_parts(p)
-	return idxs[vis], locs[vis].T
+	res = p.visible_locs()
+	return res
 
+	# idxs, locs, vis = __expand_parts(p)
+	# return idxs[vis], locs[vis].T
 
-def crops(im, xy, ratio=DEFAULT_RATIO, padding_mode="edge"):
+def crop(im, xy, ratio=DEFAULT_RATIO, padding_mode="edge"):
 	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=np.uint8)
 
 	pad_h, pad_w = crop_h // 2, crop_w // 2
 
 	padded_im = np.pad(im, [(pad_h, pad_h), (pad_w, pad_w), [0,0]], mode=padding_mode)
 
-	for i, (x, y) in enumerate(xy.T):
-		x0, y0 = x - crop_w // 2 + pad_w, y - crop_h // 2 + pad_h
-		crops[i] = padded_im[y0:y0+crop_h, x0:x0+crop_w]
+	x0, y0 = xy[0] - crop_w // 2 + pad_w, xy[1] - crop_h // 2 + pad_h
+	return padded_im[y0:y0+crop_h, x0:x0+crop_w]
 
-	return crops
+def crops(im, xy, ratio=DEFAULT_RATIO, padding_mode="edge"):
+	return np.stack([crop(im, x, y, ratio, padding_mode)
+		for (x,y) in xy.T])
 
 def visible_crops(im, p, *args, **kw):
-	idxs, locs, vis = __expand_parts(p)
-	parts = crops(asarray(im), locs[vis].T, *args, **kw)
-	res = np.zeros((len(idxs),) + parts.shape[1:], dtype=parts.dtype)
-	res[vis] = parts
+	res = p.visible_crops(*args, **kw)
 	return res
 
+	# idxs, locs, vis = __expand_parts(p)
+	# 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 = dimensions(im)
 	crop_h, crop_w = int(h * ratio), int(w * ratio)

scripts/display_from_info.py

@@ -79,13 +79,11 @@ def main(args):
 	for i in range(start, max(start, start + n_images)):
 		im, parts, label = data[i]
 
-		idxs, xy = utils.visible_part_locs(parts)
-		part_crops = utils.visible_crops(im, parts, ratio=data.ratio)
+		idxs, xy = parts.visible_locs()
+		part_crops = parts.visible_crops(ratio=data.ratio)
 		if args.rnd:
-			selected = parts[:, -1].astype(bool)
-			parts[selected, -1] = 0
-			parts[np.logical_not(selected), -1] = 1
-			action_crops = utils.visible_crops(im, parts, ratio=data.ratio)
+			parts.invert_selection()
+			action_crops = parts.visible_crops(ratio=data.ratio)
 
 		fig1 = plt.figure(figsize=(16,9))
 		ax = fig1.add_subplot(2,1,1)
@@ -103,7 +101,7 @@ def main(args):
 		plot_crops(part_crops, "Selected parts", names=crop_names)
 
 		if args.rnd:
-			plot_crops(action_crops, "Actions")
+			plot_crops(action_crops, "Actions", names=crop_names)
 
 		plt.show()
 		plt.close()