cvdatasets/nabirds/dataset/part.py

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])