Upload folder using huggingface_hub

config.json

@@ -0,0 +1,7 @@
+{
+    "model_name": "facebook/dinov3-vits16-pretrain-lvd1689m",
+    "num_classes": 145,
+    "temperature": 1.3,
+    "architecture": "DinoV3LinearMultiLinear",
+    "architecture_description": "Frozen DinoV3-ViT-S16 backbone + 3-layer MLP head (384 -> 256 -> 128 -> 145 classes)"
+  }

id2label.json

@@ -0,0 +1,147 @@
+{
+  "0": "Acantharia",
+  "1": "Acanthoica_quattrospina",
+  "2": "Akashiwo_sanguinea",
+  "3": "Alexandrium_spp",
+  "4": "Amphidinium_crassum",
+  "5": "Amphidinium_sphenoides",
+  "6": "Apedinella",
+  "7": "Askenasia",
+  "8": "Asterompalus_flabellatus",
+  "9": "Asteromphalus_sarcophagus",
+  "10": "Azadinium_caudata",
+  "11": "Bacillaria",
+  "12": "Bacteriastrum",
+  "13": "Balanion",
+  "14": "Bead",
+  "15": "Brockmanniella",
+  "16": "Calciopappus",
+  "17": "Calyptrosphaera",
+  "18": "Carchesium",
+  "19": "Cerautulina_pelagica_chain",
+  "20": "Cerautulina_pelagica_single_double",
+  "21": "Chaetoceros_curvisetus_debelis",
+  "22": "Chaetoceros_didymus",
+  "23": "Chaetoceros_didymus_single",
+  "24": "Chaetoceros_morphotype1",
+  "25": "Chaetoceros_peruvianum",
+  "26": "Chaetoceros_socialis_type",
+  "27": "Chaetoceros_spp",
+  "28": "Chaetoceros_tenuissimus",
+  "29": "Chrysochromulina",
+  "30": "Codonellopsis",
+  "31": "Corethron",
+  "32": "Coscinodiscus_granii",
+  "33": "Cryptophyta",
+  "34": "Cylindrotheca_closterium",
+  "35": "Delphineis",
+  "36": "Detonula_pumila",
+  "37": "Detritus",
+  "38": "Dictyocha_fibula",
+  "39": "Dictyocha_speculum",
+  "40": "Dinobyron",
+  "41": "Dinophysis_acuminata_complex",
+  "42": "Dinophysis_acuta",
+  "43": "Dinophysis_caudata",
+  "44": "Dinophysis_caudata_var_diegensis",
+  "45": "Dinophysis_fortii",
+  "46": "Dinophysis_tripos",
+  "47": "Diploneis_crabro",
+  "48": "Diplopsalis",
+  "49": "Ditylum_brightwellii",
+  "50": "Entomoneis",
+  "51": "Erythropsidium",
+  "52": "Eucampia_",
+  "53": "Eutintinnus",
+  "54": "Eutriptiella",
+  "55": "Faecal_pellet",
+  "56": "Favella",
+  "57": "Flagellate_clump",
+  "58": "Flagellate_heart_shape",
+  "59": "Flagellate_morphotype1",
+  "60": "Flagellate_morphotype2",
+  "61": "Gonyaulax_spinifera",
+  "62": "Gonyaulax_verior",
+  "63": "Guinardia_delicatula_chain",
+  "64": "Guinardia_delicatula_single",
+  "65": "Guinardia_delicatula_single_double",
+  "66": "Guinardia_flaccida",
+  "67": "Guinardia_striata",
+  "68": "Gymnodiniales_morphotype1",
+  "69": "Gymnodinium_catenatum",
+  "70": "Gyrodinium_falcatum",
+  "71": "Gyrodinium_spirale_type",
+  "72": "Halosphaera",
+  "73": "Haslea_wawrikae",
+  "74": "Helicotheca_tamesis",
+  "75": "Heterocapsa_azadinium",
+  "76": "Heterocapsa_rotundata",
+  "77": "Heterocapsa_type",
+  "78": "Karenia_mikimotoi",
+  "79": "Karlodinium",
+  "80": "Katodinium",
+  "81": "Laboea_strobila",
+  "82": "Lauderia_annulata",
+  "83": "Leegaardiella_sol",
+  "84": "Lessardia",
+  "85": "Mesodinium_rubrum",
+  "86": "Mesodinium_small",
+  "87": "Mesoporos",
+  "88": "Meuniera_membranacea",
+  "89": "Meuniera_membranacea_single",
+  "90": "Nauplii",
+  "91": "Navicula_transitans_var_derasa",
+  "92": "Navicula_transitrans_var_derasa_f_delicatula",
+  "93": "Odontella_mobiliensis",
+  "94": "Oxytoxum_gracile",
+  "95": "Paralia_sulcata",
+  "96": "Phaeocystis",
+  "97": "Phalachroma_rotundatum",
+  "98": "Plagiolemma_distortum",
+  "99": "Planktoniella_sol",
+  "100": "Pleurosigma",
+  "101": "Podosira_stelligera",
+  "102": "Polykrikos",
+  "103": "Proboscia_truncata",
+  "104": "Prorocentrum_cordatum",
+  "105": "Prorocentrum_cordatum_minimum",
+  "106": "Prorocentrum_dentatum",
+  "107": "Prorocentrum_gracile",
+  "108": "Prorocentrum_micans",
+  "109": "Protoperidinium",
+  "110": "Protoperidinium_bipes",
+  "111": "Protoperidinium_steinii",
+  "112": "Psammodictyon_panduriforme",
+  "113": "Pseudchattonella_farcimen_round",
+  "114": "Pseudo-nitzschia_chain_double",
+  "115": "Pseudo-nitzschia_chain_multiple",
+  "116": "Pseudo-nitzschia_single",
+  "117": "Pseudochattonella",
+  "118": "Pseudochattonella_covering_Dictyocha",
+  "119": "Pseudochattonella_farcimen_oblong",
+  "120": "Pseudosolenia_calcar-avis",
+  "121": "Pterosperma",
+  "122": "Radiolaria_lithomelissa",
+  "123": "Rotifera",
+  "124": "Scrippsiella",
+  "125": "Stenosomella",
+  "126": "Stephanopyxsis",
+  "127": "Strombidium_ciliate",
+  "128": "Thalassionema_nitzschioides_double",
+  "129": "Thalassionema_nitzschioides_multiple",
+  "130": "Thalassionema_nitzschioides_single",
+  "131": "Thalassiosira_gravida_double",
+  "132": "Thalassiosira_gravida_rotula",
+  "133": "Thalassiosira_gravida_single",
+  "134": "Tiarina_fusus",
+  "135": "Tintinnopsis",
+  "136": "Tontonia",
+  "137": "Torodinium",
+  "138": "Torodinium_teredo",
+  "139": "Tripos_furca",
+  "140": "Tripos_fusus",
+  "141": "Tripos_horridus",
+  "142": "Tripos_lineatus",
+  "143": "Tripos_muelleri",
+  "144": "Undet_small"
+}

inference.py

@@ -0,0 +1,54 @@
+import torch
+from transformers import AutoModel, AutoImageProcessor
+from model import DinoV3LinearMultiLinear
+
+def load_model(weights_path, device="cuda"):
+    """
+    Load the pre-trained classifier.
+    
+    Args:
+        weights_path: Path to the saved weights (.pt file)
+        device: Device to load model on ('cuda' or 'cpu')
+    
+    Returns:
+        model: Loaded DinoV3LinearMultiLinear model in eval mode
+        processor: Image processor for preprocessing input images
+    """
+
+    # Load config
+    import json
+    with open("config.json", "r") as f:
+        config = json.load(f)
+    # Load backbone
+    backbone = AutoModel.from_pretrained(config["model_name"])
+    hidden_size = backbone.config.hidden_size
+    # Instantiate classifier head
+    model = DinoV3LinearMultiLinear(
+        backbone=backbone,
+        num_classes=config["num_classes"],
+        hidden_size=hidden_size,
+        freeze_backbone=True
+    )
+    
+    # Load trained weights
+    model.load_state_dict(torch.load(weights_path, map_location=device)["model_state_dict"])
+    model.to(device)
+    model.eval()
+    
+    # Load image processor
+    processor = AutoImageProcessor.from_pretrained(config["model_name"])
+
+    # Load labels
+    with open("id2label.json", "r") as f:
+        id2label = json.load(f)
+    
+    return model, processor, id2label
+    
+
+def probs_to_labels(probs, id2label):
+    """
+    Convert probability distribution to labels.
+    """
+    predicted_indices = probs.argmax(dim=1)
+    predicted_labels = [id2label[str(idx.item())] for idx in predicted_indices]
+    return predicted_labels

model.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModel
+
+class DinoV3LinearMultiLinear(nn.Module):
+    def __init__(self, backbone: AutoModel, hidden_size: int, num_classes: int, freeze_backbone: bool = True):
+        super().__init__()
+        self.backbone = backbone
+        self.num_classes = num_classes
+        if freeze_backbone:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+            self.backbone.eval()
+        # three linear layers like in the original syke-pic model
+        # hidden size -> 256 -> 128 -> num_classes
+        self.linear1 = nn.Linear(hidden_size, 256)
+        self.linear2 = nn.Linear(256, 128)
+        self.linear3 = nn.Linear(128, self.num_classes)
+
+    def print_num_trainable_parameters(self):
+        print(f"Number of trainable parameters: {sum(p.numel() for p in self.parameters() if p.requires_grad)}")
+
+    def forward(self, pixel_values):
+        outputs = self.backbone(pixel_values=pixel_values)
+        last_hidden = outputs.last_hidden_state
+        cls = last_hidden[:, 0]
+        logits = self.linear3(self.linear2(self.linear1(cls)))
+        return logits
+
+    def predict(self, pixel_values, temperature=1.3):
+        """
+        Generate probability predictions for a batch of images.
+        
+        Args:
+            pixel_values: Preprocessed image tensor (batch_size, 3, H, W)
+            temperature: Temperature for softmax calibration (default 1.3)
+        
+        Returns:
+            probs: Probability distribution over classes (shape: [batch_size, num_classes])
+        """
+        logits = self.forward(pixel_values)
+        probs = torch.softmax(logits / temperature, dim=1)
+        return probs

readme.md

@@ -0,0 +1,19 @@
+---
+license: mit
+language:
+  - en
+tags:
+  - image-classification
+  - plankton
+  - dinov3
+  - biology
+  - marine
+datasets:
+  - ifcb
+---
+
+# DINO Plankton Classifier
+The model is trained on PML IFCB data consisting of 145 plankton classes.
+
+# Inference
+Use the provided inference script. See example in `demo.ipynb` on predicting the classes for two synthetic samples.

requirements.txt

@@ -0,0 +1,16 @@
+accelerate==1.10.1
+bitsandbytes==0.48.2
+datasets==2.21.0
+diffusers==0.35.1
+evaluate==0.4.5
+fastapi==0.116.1
+ffmpy==0.6.1
+gradio==5.45.0
+gradio_client==1.13.0
+safehttpx==0.1.6
+safetensors==0.6.2
+tokenizers==0.22.0
+torch==2.8.0+cu126
+torchaudio==2.8.0+cu126
+torchvision==0.23.0+cu126
+transformers==4.57.0.dev0

weights.pt

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