Added check_dataset_output

Author: sadda

wildlife_tools/features/deep.py

@@ -3,6 +3,7 @@
 from transformers import CLIPModel, CLIPProcessor
 from typing import Optional
 from ..data import FeatureDataset, ImageDataset
+from ..tools import check_dataset_output
 
 
 class DeepFeatures:
@@ -43,7 +44,8 @@ def __call__(self, dataset: ImageDataset) -> FeatureDataset:
 
         self.model = self.model.to(self.device)
         self.model = self.model.eval()
-
+        
+        check_dataset_output(dataset, check_label=False)
         loader = torch.utils.data.DataLoader(
             dataset,
             num_workers=self.num_workers,
@@ -119,6 +121,7 @@ def __call__(self, dataset: ImageDataset) -> FeatureDataset:
         # TODO: this is hacky
         dataset.transforms = None  # Reset transforms.
 
+        check_dataset_output(dataset, check_label=False)
         loader = torch.utils.data.DataLoader(
             dataset,
             num_workers=self.num_workers,
@@ -127,7 +130,7 @@ def __call__(self, dataset: ImageDataset) -> FeatureDataset:
             collate_fn=lambda x: x,
         )
         outputs = []
-        for image in tqdm(loader, mininterval=1, ncols=100):
+        for image, _ in tqdm(loader, mininterval=1, ncols=100):
             with torch.no_grad():
                 output = self.model(self.transform(image).to(self.device)).pooler_output
                 outputs.append(output.cpu())

wildlife_tools/features/local.py

@@ -6,6 +6,7 @@
 from tqdm import tqdm
 
 from ..data import FeatureDataset, ImageDataset
+from ..tools import check_dataset_output
 from .gluefactory_fix import extract_single_image_fix  # https://github.com/cvg/glue-factory/pull/50
 
 
@@ -53,6 +54,7 @@ def __call__(self, dataset: ImageDataset) -> FeatureDataset:
         """
 
         features = []
+        check_dataset_output(dataset, check_label=False)
         loader = torch.utils.data.DataLoader(
             dataset,
             num_workers=self.num_workers,

wildlife_tools/features/memory.py

@@ -1,6 +1,6 @@
 from tqdm import tqdm
-
 from ..data import FeatureDataset, ImageDataset
+from ..tools import check_dataset_output
 
 
 class DataToMemory:
@@ -17,6 +17,7 @@ def __call__(self, dataset: ImageDataset):
         """Loads data from input dataset into array and returns them as a new FeatureDataset."""
 
         features = []
+        check_dataset_output(dataset, check_label=False)
         for x, _ in tqdm(dataset, mininterval=1, ncols=100):
             features.append(x)
         return FeatureDataset(

wildlife_tools/tools.py

@@ -8,6 +8,23 @@
 from PIL import Image
 
 
+def check_dataset_output(dataset, check_label=False):
+    output = dataset[0]
+    if not isinstance(output, tuple):
+        raise ValueError('''
+            Calling `dataset[0]` must returned a tuple.
+            Try to use `load_label=True` when creating the dataset.
+            ''')
+    label = output[1]
+    if check_label and (isinstance(label, str) or isinstance(label, np.str_)):
+        raise ValueError('''
+            Calling `dataset[0]` must returned a tuple,
+            where the second part (label) is an integer.
+            If you used the WildlifeDataset from wildlife-datasets,
+            try to use `factorize_label=True` when creating the dataset.
+            ''')
+
+
 def frame_image(img, frame_width, color=(255, 0, 0)):
     b = frame_width
     ny, nx = img.shape[0], img.shape[1]

wildlife_tools/train/trainer.py

@@ -6,6 +6,7 @@
 from tqdm import tqdm
 from typing import Optional, Callable
 from ..data import ImageDataset
+from ..tools import check_dataset_output
 
 
 def set_seed(seed=0, device="cuda"):
@@ -94,6 +95,7 @@ def __init__(
         accumulation_steps: int = 1,
         epoch_callback: Optional[Callable] = None,
     ):
+        check_dataset_output(dataset, check_label=True)
         self.dataset = dataset
         self.model = model.to(device)
         self.objective = objective.to(device)