Remove species_ prefix dependency from image path

tests/test_image_file_handling.py

@@ -308,7 +308,7 @@ def test_train_with_md_labels(mocker, ena24_dataset_setup):
     assert set(["x1", "y1", "x2", "y2"]).issubset(config.labels.columns)
 
     # MD bounds are relative
-    assert config.labels.loc[0, "x1"] < 1.0
+    assert config.labels.iloc[0]["x1"] < 1.0
 
     data = ImageClassificationDataModule(
         data_dir=config.data_dir,
@@ -320,15 +320,15 @@ def test_train_with_md_labels(mocker, ena24_dataset_setup):
         crop_images=config.crop_images,
     )
 
-    assert data.annotations.loc[0, "x1"] == 1920 * 0.35  # 672
-    assert data.annotations.loc[0, "y1"] == 1080 * 0.35  # 378
+    assert data.annotations.iloc[0]["x1"] == 1920 * 0.35  # 672
+    assert data.annotations.iloc[0]["y1"] == 1080 * 0.35  # 378
 
     # some buffer for floating point precision
     assert (
-        abs(data.annotations.loc[0, "x2"] - ((1920 * 0.35) + (1920 * 0.3))) <= 1
+        abs(data.annotations.iloc[0]["x2"] - ((1920 * 0.35) + (1920 * 0.3))) <= 1
     )  # 672 + 576 = 1248
     assert (
-        abs(data.annotations.loc[0, "y2"] - ((1080 * 0.35) + (1080 * 0.3))) <= 1
+        abs(data.annotations.iloc[0]["y2"] - ((1080 * 0.35) + (1080 * 0.3))) <= 1
     )  # 378 + 324 = 702
 
     # make sure bounding boxes are absolute

zamba/images/config.py

@@ -6,6 +6,7 @@
 
 import appdirs
 import pandas as pd
+from sklearn.preprocessing import LabelEncoder
 import torch
 from loguru import logger
 from pydantic import DirectoryPath, FilePath, root_validator, validator
@@ -403,32 +404,26 @@ def preprocess_labels(cls, values):
         labels = values["labels"]
 
         # lowercase to facilitate subset checking
-        labels["label"] = labels.label.str.lower()
-
-        # one hot encoding
-        labels = pd.get_dummies(labels.rename(columns={"label": "species"}), columns=["species"])
+        labels["class_name"] = labels.label.str.lower()
 
         # We validate that all the images exist prior to this, so once this assembles the set of classes,
         # we should have at least one example of each label and don't need to worry about filtering out classes
         # with missing examples.
-        species_columns = labels.columns[labels.columns.str.contains("species_")]
-        values["species_in_label_order"] = species_columns.to_list()
+        encoder = LabelEncoder()
+        labels["label"] = encoder.fit_transform(labels["class_name"])
+        values["species_in_label_order"] = encoder.classes_.tolist()
 
-        indices = (
-            labels[species_columns].idxmax(axis=1).apply(lambda x: species_columns.get_loc(x))
-        )
-
-        labels["label"] = indices
+        # one hot encoding
+        one_hot = pd.get_dummies(labels["class_name"])
+        labels = labels.assign(**one_hot)
 
         # if no "split" column, set up train, val, and test split
         if "split" not in labels.columns:
             make_split(labels, values)
 
-        values["labels"] = labels.reset_index()
+        values["labels"] = labels
 
-        example_species = [
-            species.replace("species_", "") for species in values["species_in_label_order"][:3]
-        ]
+        example_species = values["species_in_label_order"][:3]
         logger.info(
             f"Labels preprocessed. {len(values['species_in_label_order'])} species found: {example_species}..."
         )

zamba/images/manager.py

@@ -9,6 +9,7 @@
 
 import git
 import mlflow
+import numpy as np
 import pandas as pd
 import pytorch_lightning as pl
 import torch
@@ -39,11 +40,8 @@
 from zamba.pytorch.transforms import resize_and_pad
 
 
-def get_weights(split):
-    labels_df = split.filter(like="species_")
-    y_array = pd.from_dummies(labels_df).values.flatten()
-    classes = labels_df.columns.values
-    class_weights = compute_class_weight("balanced", classes=classes, y=y_array)
+def get_weights(split, all_labels):
+    class_weights = compute_class_weight("balanced", classes=all_labels, y=split.label)
     return torch.tensor(class_weights).to(torch.float32)
 
 
@@ -250,7 +248,9 @@
 
     loss_fn = torch.nn.CrossEntropyLoss()
     if config.weighted_loss is True:
-        loss_fn = torch.nn.CrossEntropyLoss(weight=get_weights(data.annotations), reduction="mean")
+        loss_fn = torch.nn.CrossEntropyLoss(
+            weight=get_weights(data.annotations, np.unique(config.labels.label)), reduction="mean"
+        )
 
     # Calculate number of training batches
     num_training_batches = len(data.train_dataloader())

zamba/models/config.py

@@ -671,9 +671,15 @@ def make_split(labels, values):
 
         # check we have at least as many videos per species as we have splits
         # labels are OHE at this point
-        num_videos_per_species = labels.filter(regex="species_").sum().to_dict()
+        if "species_in_label_order" in values:
+            num_videos_per_species = labels.class_name.value_counts().to_dict()
+            species = values["species_in_label_order"]
+        else:
+            num_videos_per_species = labels.filter(regex="species_").sum().to_dict()
+            species = labels.filter(regex="species_").columns
+
         too_few = {
-            k.split("species_", 1)[1]: v
+            k.removeprefix("species_"): v
             for k, v in num_videos_per_species.items()
             if 0 < v < len(expected_splits)
         }
@@ -683,7 +689,7 @@ def make_split(labels, values):
                 f"Not all species have enough media files to allocate into the following splits: {', '.join(expected_splits)}. A minimum of {len(expected_splits)} media files per label is required. Found the following counts: {too_few}. Either remove these labels or add more images/videos."
             )
 
-        for c in labels.filter(regex="species_").columns:
+        for c in species:
             species_df = labels[labels[c] > 0]
 
             if len(species_df):