Browse Source

added bounding box loading and cropping

Dimitri Korsch 6 years ago
parent
commit
ff97a2e891
3 changed files with 51 additions and 6 deletions
  1. 5 2
      Makefile
  2. 22 2
      nabirds/annotations.py
  3. 24 2
      nabirds/dataset/__init__.py

+ 5 - 2
Makefile

@@ -1,5 +1,5 @@
 install:
-	python setup.py install
+	pip install . --no-deps --upgrade
 
 build:
 	python setup.py build
@@ -7,5 +7,8 @@ build:
 deploy:
 	python setup.py sdist upload -r pypi
 
+test_deploy:
+	python setup.py sdist upload -r pypitest
+
 get_version:
-	@python -c "import nabirds; print('v{}'.format(nabirds.__version__))"
+	@python -c "import cvargparse; print('v{}'.format(cvargparse.__version__))"

+ 22 - 2
nabirds/annotations.py

@@ -86,10 +86,10 @@ class BaseAnnotations(abc.ABC):
 		return self.image_path(fname)
 
 	def label(self, uuid):
-		return self.labels[self.uuid_to_idx[uuid]]
+		return self.labels[self.uuid_to_idx[uuid]].copy()
 
 	def parts(self, uuid):
-		return self.part_locs[self.uuid_to_idx[uuid]]
+		return self.part_locs[self.uuid_to_idx[uuid]].copy()
 
 
 	def _uuids(self, split):
@@ -134,8 +134,11 @@ class CUB_Annotations(BaseAnnotations):
 			images_file="images.txt",
 			labels_file="labels.txt",
 			split_file="tr_ID.txt",
+			bounding_boxes="bounding_boxes.txt",
+			bounding_box_dtype=np.dtype([(v, np.int32) for v in "xywh"]),
 			parts_file=join("parts", "part_locs.txt"),
 			part_names_file=join("parts", "parts.txt"),
+
 		)
 
 		info.structure = [
@@ -144,6 +147,7 @@ class CUB_Annotations(BaseAnnotations):
 			[info.split_file, "_split"],
 			[info.parts_file, "_part_locs"],
 			[info.part_names_file, "_part_names"],
+			[info.bounding_boxes, "_bounding_boxes"],
 		]
 		return info
 
@@ -163,4 +167,20 @@ class CUB_Annotations(BaseAnnotations):
 		# set part idxs from 1-idxs to 0-idxs
 		self.part_locs[..., 0] -= 1
 
+		self._load_bounding_boxes()
+
+	def _load_bounding_boxes(self):
+		assert self._bounding_boxes is not None, "Bouding boxes were not loaded!"
+
+		uuid_to_bbox = {}
+		for content in [i.split() for i in self._bounding_boxes]:
+			uuid, bbox = content[0], content[1:]
+			uuid_to_bbox[uuid] = [float(i) for i in bbox]
+
+		self.bounding_boxes = np.array(
+			[tuple(uuid_to_bbox[uuid]) for uuid in self.uuids],
+			dtype=self.meta.bounding_box_dtype)
+
 
+	def bounding_box(self, uuid):
+		return self.bounding_boxes[self.uuid_to_idx[uuid]].copy()

+ 24 - 2
nabirds/dataset/__init__.py

@@ -3,10 +3,25 @@ import numpy as np
 
 
 class Dataset(object):
-	def __init__(self, uuids, annotations):
+	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)
@@ -15,11 +30,18 @@ class Dataset(object):
 		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)