cvdatasets/nabirds/dataset/__init__.py

from imageio import imread
import numpy as np


class Dataset(object):
	def __init__(self, uuids, annotations, crop_to_bb=False, crop_uniform=False):
		super(Dataset, self).__init__()
		self.uuids = uuids
		self._annot = annotations
		self.crop_to_bb = crop_to_bb
		self.crop_uniform = crop_uniform

	def bounding_box(self, i):
		bbox = self._get("bounding_box", i)
		x,y,w,h = [bbox[attr] for attr in "xywh"]
		if self.crop_uniform:
			x0 = x + w//2
			y0 = y + h//2

			crop_size = max(w//2, h//2)

			x,y = max(x0 - crop_size, 0), max(y0 - crop_size, 0)
			w = h = crop_size * 2
		return x,y,w,h

	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, mode="RGB"):
		methods = ["image", "parts", "label"]
		im_path, parts, label = [self._get(m, i) for m in methods]

		im = imread(im_path, pilmode=mode)

		if self.crop_to_bb:
			x,y,w,h = self.bounding_box(i)
			im = im[y:y+h, x:x+w]
			parts[:, 1] -= x
			parts[:, 2] -= y

		h,w,c = im.shape
		# fit to the dimensions of the image
		parts[:, 1] = np.minimum(parts[:, 1], w - 1)
		parts[:, 2] = np.minimum(parts[:, 2], h - 1)
		return im, parts, label

	__getitem__  = get_example