Use remote URLs for new models. Update formatting.

RolnickLab · Aug 15, 2024 · f5d5cc1 · f5d5cc1
1 parent 4e9cf3a
commit f5d5cc1
Showing 1 changed file with 12 additions and 13 deletions.
diff --git a/trapdata/ml/models/classification.py b/trapdata/ml/models/classification.py
@@ -1,5 +1,6 @@
 import timm
 import torch
+import torch.utils.data
 import torchvision
 
 from trapdata import constants, logger
@@ -105,6 +106,7 @@ def forward(self, x):
 
  return x
 
+
 class Resnet50Classifier_Turing(InferenceBaseClass):
  # function to run the Turing models
  logger.info("KG: Resnet50Classifier_Turing")
@@ -144,6 +146,7 @@ def post_process_batch(self, output):
  logger.debug(f"Post-processing result batch: {result}")
  return result
 
+
 class Resnet50Classifier(InferenceBaseClass):
  input_size = 300
 
@@ -299,37 +302,33 @@ class QuebecVermontMothSpeciesClassifierMixedResolution(
  "quebec-vermont_moth-category-map_19Jan2023.json"
  )
 
-class TuringCostaRicaSpeciesClassifier(
- SpeciesClassifier, Resnet50Classifier_Turing
-):
+
+class TuringCostaRicaSpeciesClassifier(SpeciesClassifier, Resnet50Classifier_Turing):
  name = "Turing Costa Rica Species Classifier"
- description = (
- "Trained on 4th June 2024 by Turing team using Resnet50 model."
- )
+ description = "Trained on 4th June 2024 by Turing team using Resnet50 model."
  weights_path = (
+ "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
  "turing-costarica_v03_resnet50_2024-06-04-16-17_state.pt"
  )
  labels_path = (
+ "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
  "03_costarica_data_category_map.json"
  )
 
 
-class TuringUKSpeciesClassifier(
- SpeciesClassifier, Resnet50Classifier_Turing
-):
+class TuringUKSpeciesClassifier(SpeciesClassifier, Resnet50Classifier_Turing):
  name = "Turing UK Species Classifier"
- description = (
- "Trained on 13th May 2024 by Turing team using Resnet50 model."
- )
+ description = "Trained on 13th May 2024 by Turing team using Resnet50 model."
  weights_path = (
+ "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
  "turing-uk_v03_resnet50_2024-05-13-10-03_state.pt"
  )
  labels_path = (
+ "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
  "03_uk_data_category_map.json"
  )
 
 
-
 class UKDenmarkMothSpeciesClassifierMixedResolution(
  SpeciesClassifier, Resnet50ClassifierLowRes
 ):