瀏覽代碼

refactored annotation type definitions and handling

Dimitri Korsch 5 年之前
父節點
當前提交
a3fce08516
共有 2 個文件被更改,包括 88 次插入36 次删除
  1. 20 36
      cvdatasets/annotations/__init__.py
  2. 68 0
      cvdatasets/annotations/annotation_types.py

+ 20 - 36
cvdatasets/annotations/__init__.py

@@ -1,3 +1,4 @@
+from cvdatasets.annotations.annotation_types import AnnotationType
 from cvdatasets.annotations.impl.birdsnap import BSNAP_Annotations
 from cvdatasets.annotations.impl.cars import CARS_Annotations
 from cvdatasets.annotations.impl.cub import CUB_Annotations
@@ -11,43 +12,26 @@ from cvdatasets.annotations.impl.inat import INAT20_Annotations
 from cvdatasets.annotations.impl.nab import NAB_Annotations
 from cvdatasets.annotations.impl.tigers import TIGERS_Annotations
 
+
 from cvdatasets.annotations.base import BaseAnnotations
 from cvdatasets.annotations.base.bbox_mixin import BBoxMixin
 from cvdatasets.annotations.base.parts_mixin import PartsMixin
 
-from cvargparse.utils import BaseChoiceType
-from functools import partial
-
-class AnnotationType(BaseChoiceType):
-	CUB200 = CUB_Annotations
-	CUB200_2FOLD = partial(CUB_Annotations)
-	CUB200_GOOGLE = partial(CUB_Annotations)
-	CUB200_GOOGLE_SEM = partial(CUB_Annotations)
-	BIRDSNAP = BSNAP_Annotations
-	NAB = NAB_Annotations
-	CARS = CARS_Annotations
-	DOGS = DOGS_Annotations
-	FLOWERS = FLOWERS_Annotations
-	HED = HED_Annotations
-	TIGERS = TIGERS_Annotations
-	TIGERS_TEST = partial(TIGERS_Annotations)
-
-	IMAGENET = INET_Annotations
-
-	INAT18 = INAT18_Annotations
-
-	INAT19 = INAT19_Annotations
-	INAT19_MINI = partial(INAT19_Annotations)
-	INAT19_TEST = partial(INAT19_Annotations)
-
-	INAT20 = INAT20_Annotations
-	INAT20_TEST = partial(INAT20_Annotations)
-	INAT20_IN_CLASS = partial(INAT20_Annotations)
-	INAT20_OUT_CLASS = partial(INAT20_Annotations)
-	INAT20_NOISY_IN_CLASS = partial(INAT20_Annotations)
-	INAT20_U_OUT_CLASS = partial(INAT20_Annotations)
-	INAT20_NOISY_OUT_CLASS = partial(INAT20_Annotations)
-
-	IMAGENET_TOP_INAT20 = partial(INET_Annotations)
-
-	Default = CUB200
+__all__ = [
+	"AnnotationType",
+	"BaseAnnotations",
+	"BBoxMixin",
+	"BSNAP_Annotations",
+	"CARS_Annotations",
+	"CUB_Annotations",
+	"DOGS_Annotations",
+	"FLOWERS_Annotations",
+	"HED_Annotations",
+	"INAT18_Annotations",
+	"INAT19_Annotations",
+	"INAT20_Annotations",
+	"INET_Annotations",
+	"NAB_Annotations",
+	"PartsMixin",
+	"TIGERS_Annotations",
+]

+ 68 - 0
cvdatasets/annotations/annotation_types.py

@@ -0,0 +1,68 @@
+from cvdatasets.annotations.impl.birdsnap import BSNAP_Annotations
+from cvdatasets.annotations.impl.cars import CARS_Annotations
+from cvdatasets.annotations.impl.cub import CUB_Annotations
+from cvdatasets.annotations.impl.dogs import DOGS_Annotations
+from cvdatasets.annotations.impl.flowers import FLOWERS_Annotations
+from cvdatasets.annotations.impl.hed import HED_Annotations
+from cvdatasets.annotations.impl.imagenet import INET_Annotations
+from cvdatasets.annotations.impl.inat import INAT18_Annotations
+from cvdatasets.annotations.impl.inat import INAT19_Annotations
+from cvdatasets.annotations.impl.inat import INAT20_Annotations
+from cvdatasets.annotations.impl.nab import NAB_Annotations
+from cvdatasets.annotations.impl.tigers import TIGERS_Annotations
+
+from cvargparse.utils import BaseChoiceType
+from functools import partial
+
+class AnnotationType(BaseChoiceType):
+	IMAGENET = INET_Annotations
+
+	CUB200 = CUB_Annotations
+	BIRDSNAP = BSNAP_Annotations
+	NAB = NAB_Annotations
+
+	CARS = CARS_Annotations
+	DOGS = DOGS_Annotations
+
+	FLOWERS = FLOWERS_Annotations
+
+	HED = HED_Annotations
+	TIGERS = TIGERS_Annotations
+
+	INAT18 = INAT18_Annotations
+	INAT19 = INAT19_Annotations
+	INAT20 = INAT20_Annotations
+
+	Default = CUB200
+
+	@classmethod
+	def phony(cls, key):
+		""" returns for a key a list of datasets,
+			that use the same annotation class """
+
+		return {
+			cls.CUB200 : [ "CUB200_2FOLD", "CUB200_GOOGLE", "CUB200_GOOGLE_SEM" ],
+			cls.TIGERS : [ "TIGERS_TEST" ],
+			cls.INAT19 : [ "INAT19_TEST", "INAT19_MINI" ],
+			cls.INAT20 : [ "INAT20_TEST",
+				"INAT20_IN_CLASS",
+				"INAT20_OUT_CLASS",
+				"INAT20_NOISY_IN_CLASS",
+				"INAT20_NOISY_OUT_CLASS",
+				"INAT20_U_IN_CLASS",
+				"INAT20_U_OUT_CLASS",
+			],
+			cls.IMAGENET : [ "IMAGENET_TOP_INAT20" ],
+		}.get(key, [])
+
+	@classmethod
+	def as_choices(cls, add_phony=True):
+		choices = super(AnnotationType, cls).as_choices()
+		if not add_phony:
+			return choices
+
+		for key in cls:
+			for phony in cls.phony(key):
+				choices[phony.lower()] = choices[key.name.lower()]
+
+		return choices