cvdatasets/nabirds/dataset/part.py

import numpy as np
from contextlib import contextmanager
from abc import ABC, abstractmethod, abstractproperty

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 = [BasePart.new(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 offset(self, dx, dy):
		for p in self._parts:
			p.x += dx
			p.y += dy

	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 BasePart(ABC):
	def __init__(self, image, annotation):
		super(BasePart, self).__init__()
		self.image = image
		self.read_annotation(annotation)

	@staticmethod
	def new(image, annotation):
		if len(annotation) == 4:
			return LocationPart(image, annotation)
		elif len(annotation) == 5:
			return BBoxPart(image, annotation)
		else:
			raise ValueError("Unknown part annotation format: {}".format(annotation))


	@abstractmethod
	def read_annotation(self, annotation):
		raise NotImplementedError

	@property
	def is_visible(self):
		return 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 c(self):
		h, w, c = utils.dimensions(self.image)
		return c

	@abstractmethod
	def crop(self, ratio=None, padding_mode="edge"):
		raise NotImplementedError


class LocationPart(BasePart):

	def read_annotation(self, annotation):
		# here x,y are the center of the part
		self._id, self.x, self.y, self.is_visible = annotation
		self._ratio = None

	@abstractmethod
	def crop(self, padding_mode="edge", *args, **kwargs):
		raise NotImplementedError

class LocationPart(BasePart):

	def read_annotation(self, annotation):
		# here x,y are the center of the part
		self._id, self.x, self.y, self.is_visible = annotation
		self._ratio = None

	def crop(self, image, ratio=None, padding_mode="edge", *args, **kwargs):
		ratio = ratio or self._ratio
		_h, _w, c = utils.dimensions(image)
		w, h = int(_w * ratio), int(_h * ratio)

		if not self.is_visible:
			return np.zeros((h, w, c), dtype=np.uint8)
		else:
			return utils.crop(image, self.xy, w, h,
				padding_mode, is_location=True)

class BBoxPart(BasePart):

	def read_annotation(self, annotation):
		# 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

	def crop(self, padding_mode="edge", *args, **kwargs):
		return utils.crop(self.image, self.xy, self.w, self.h,
			padding_mode, is_location=False)