added part reading support for inat annotations

nabirds/annotations/base.py

@@ -138,15 +138,17 @@ class BaseAnnotations(abc.ABC):
 		self.uuids, self.images = map(np.array, zip(*uuid_fnames))
 		self.uuid_to_idx = {uuid: i for i, uuid in enumerate(self.uuids)}
 
-	def _load_parts(self):
-		assert self._part_locs is not None, "Part locations were not loaded!"
+	def _load_parts(self, idx_offset=0):
+		assert self.has_parts, "Part locations were not loaded!"
 		# this part is quite slow... TODO: some runtime improvements?
-		uuid_to_parts = defaultdict(list)
+		idx_to_parts = defaultdict(list)
 		for content in [i.split() for i in self._part_locs]:
-			uuid_to_parts[content[0]].append([float(i) for i in content[1:]])
+			idx = int(content[0]) - idx_offset
+			idx_to_parts[idx].append([float(c) for c in content[1:]])
 
-		uuid_to_parts = dict(uuid_to_parts)
-		self.part_locs = np.stack([uuid_to_parts[uuid] for uuid in self.uuids]).astype(int)
+		idx_to_parts = dict(idx_to_parts)
+		self.part_locs = np.stack([
+			idx_to_parts[self.uuid_to_idx[uuid]] for uuid in self.uuids]).astype(int)
 
 		if hasattr(self, "_part_names") and self._part_names is not None:
 			self._load_part_names()

nabirds/annotations/cars.py

@@ -48,7 +48,7 @@ class CARS_Annotations(BaseAnnotations):
 
 		# load only if present
 		if self.has_parts:
-			super(CARS_Annotations, self)._load_parts()
+			super(CARS_Annotations, self)._load_parts(idx_offset=1)
 
 		self._load_bounding_boxes()
 

nabirds/annotations/inat.py

@@ -20,14 +20,17 @@ class INAT19_Annotations(BaseAnnotations):
 			# train_content="train2019.json",
 
 			# fake bounding boxes: the whole image
-			parts_file=join("parts", "part_locs.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 = [
 			[info.content, "_content"],
 			[info.val_content, "_val_content"],
 			[info.parts_file, "_part_locs"],
+			[info.part_names_file, "_part_names"],
 		]
 		return info
 
@@ -58,9 +61,7 @@ class INAT19_Annotations(BaseAnnotations):
 
 		# load only if present
 		if self.has_parts:
-			super(INAT19_Annotations, self)._load_parts()
-			# set part idxs from 1-idxs to 0-idxs
-			self.part_locs[..., 0] -= 1
+			super(INAT19_Annotations, self)._load_parts(idx_offset=1)
 
 		self._load_bounding_boxes()