cvdatasets/nabirds/dataset/mixins/reading.py

import numpy as np

from os.path import join

from . import BaseMixin
from ..image import ImageWrapper

class AnnotationsReadMixin(BaseMixin):

	def __init__(self, uuids, annotations, mode="RGB"):
		super(AnnotationsReadMixin, self).__init__()
		self.uuids = uuids
		self._annot = annotations
		self.mode = mode

	def __len__(self):
		return len(self.uuids)

	def _get(self, method, i):
		return getattr(self._annot, method)(self.uuids[i])

	def get_example(self, i):
		res = super(AnnotationsReadMixin, self).get_example(i)
		# if the super class returns something, then the class inheritance is wrong
		assert res is None, "AnnotationsReadMixin should be the last class in the hierarchy!"

		methods = ["image", "parts", "label"]
		im_path, parts, label = [self._get(m, i) for m in methods]

		return ImageWrapper(im_path, int(label), parts, mode=self.mode)

	@property
	def labels(self):
		return np.array([self._get("label", i) for i in range(len(self))])


class ImageListReadingMixin(BaseMixin):

	def __init__(self, pairs, root="."):
		super(ImageListReadingMixin, self).__init__()
		with open(pairs) as f:
			self._pairs = [line.strip().split() for line in f]

		assert all([len(pair) == 2 for pair in self._pairs]), \
			"Invalid format of the pairs file!"

		self._root = root

	def __len__(self):
		return len(self._pairs)

	def get_example(self, i):
		im_file, label = self._pairs[i]
		im_path = join(self._root, im_file)

		return ImageWrapper(im_path, int(label))

	@property
	def labels(self):
		return np.array([label for (_, label) in self._pairs])