feat: update moth/non-moth classifier to latest version

RolnickLab · Aug 29, 2024 · bc6a2d5 · bc6a2d5
1 parent 4a08564
commit bc6a2d5
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/trapdata/ml/models/classification.py b/trapdata/ml/models/classification.py
@@ -217,7 +217,7 @@ def get_model(self):
         return model
 
 
-class BinaryClassifier(EfficientNetClassifier):
+class BinaryClassifier(Resnet50ClassifierLowRes):
     stage = 2
     type = "binary_classification"
     positive_binary_label: str = constants.POSITIVE_BINARY_LABEL
@@ -248,7 +248,7 @@ def save_results(self, object_ids, batch_output, *args, **kwargs):
         save_classified_objects(self.db_path, object_ids, classified_objects_data)
 
 
-class MothNonMothClassifier(BinaryClassifier):
+class MothNonMothClassifier2022(EfficientNetClassifier, BinaryClassifier):
     name = "Moth / Non-Moth Classifier"
     description = "Trained on May 6, 2022"
     weights_path = "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/moth-nonmoth-effv2b3_20220506_061527_30.pth"
@@ -257,6 +257,21 @@ class MothNonMothClassifier(BinaryClassifier):
     negative_binary_label = "nonmoth"
 
 
+class MothNonMothClassifier(BinaryClassifier):
+    name = "Moth / Non-Moth Classifier"
+    description = "Trained on April 17, 2024"
+    weights_path = (
+        "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
+        "moth-nonmoth_resnet50_20240417_b4fe3efe.pth"
+    )
+    labels_path = (
+        "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
+        "05-moth-nonmoth_category_map.json"
+    )
+    positive_binary_label = "moth"
+    negative_binary_label = "nonmoth"
+
+
 class SpeciesClassifier(InferenceBaseClass):
     stage = 4
     type = "fine_grained_classifier"