Add andrena classifier

classifier-andrena-classmapping.txt

@@ -0,0 +1,50 @@
+11172350
+11178877
+1356887
+1356889
+1356976
+1356981
+1357001
+1357040
+1357066
+1357077
+1357111
+1357120
+1357156
+1357176
+1357260
+1357265
+1357279
+1357310
+1357338
+1357364
+9740714
+1357406
+1357486
+1357531
+1357545
+1357548
+1357652
+1357663
+1357665
+1357703
+1357773
+1357835
+1357872
+1357893
+1357907
+1357937
+1357950
+1358038
+1358070
+1358080
+1358088
+1358128
+1358155
+1358165
+1358224
+1358225
+1358232
+1358240
+8460383
+9805254

classifier-andrena.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c891ce12e38af92df1a0bdf75a43db93132f14e4eba7e1b52f02c32fc5091492
+size 1063717

pytorch-to-onnx.py

@@ -0,0 +1,53 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "onnx",
+#     "onnxscript",
+#     "rich",
+#     "timm",
+#     "torch",
+# ]
+# ///
+import sys
+from pathlib import Path
+
+import timm
+import torch
+from rich import print
+
+if len(sys.argv) < 3:
+    print(
+        """Usage: uv run pytorch-to-onnx.py BASE_MODEL PATH_TO_PTH PATH_TO_CLASSMAPPING_TXT IMAGE_INPUT_SIZE
+
+BASE_MODEL: The name of the base model to use (e.g., resnet50.a1_in1k). Should be available in the timm library.
+PATH_TO_PTH: can be a .tar containing a .pth folder, or a .pth file
+PATH_TO_CLASSMAPPING_TXT: A text file containing the class names, one per line. Only used to determine the number of classes for the model.
+IMAGE_INPUT_SIZE: The size of the input images (e.g., 224). The model will expect input tensors of shape [1, 3, IMAGE_INPUT_SIZE, IMAGE_INPUT_SIZE].
+        
+        """
+    )
+    sys.exit()
+
+classmapping = list(Path(sys.argv[3]).read_text().splitlines())
+
+base_model = timm.create_model(
+    sys.argv[1], pretrained=True, num_classes=len(classmapping)
+)
+
+filename = Path(sys.argv[2])
+
+
+image_input_size = int(sys.argv[4])
+
+
+state = torch.load(filename, map_location=torch.device("cpu"), weights_only=False)
+base_model.load_state_dict(state["state_dict"])
+
+base_model = torch.nn.Sequential(base_model, torch.nn.Softmax(dim=1))
+base_model.eval()
+
+torch.onnx.export(
+    base_model,
+    args=(torch.zeros([1, 3, image_input_size, image_input_size]),),
+    f=filename.with_suffix(".onnx"),
+)