Add initial submission files

data.py

@@ -0,0 +1,69 @@
+from pathlib import Path
+
+import pandas as pd
+import pytorch_lightning as pl
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import v2
+from transformers import AutoImageProcessor, AutoModel
+
+
+class TransformDino(v2.Transform):
+    def __init__(self, model_name="facebook/dinov2-base"):
+        super().__init__()
+        self.processor = AutoImageProcessor.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+
+    def forward(self, batch):
+        model_inputs = self.processor(images=batch["features"], return_tensors="pt")
+        with torch.no_grad():
+            outputs = self.model(**model_inputs)
+            last_hidden_states = outputs.last_hidden_state
+        # extract the cls token
+        batch["features"] = last_hidden_states[:, 0]
+        return batch
+
+
+class ImageDataset(Dataset):
+    def __init__(self, metadata_path, images_root_path):
+        self.metadata_path = metadata_path
+        self.metadata = pd.read_csv(metadata_path)
+        self.images_root_path = images_root_path
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        row = self.metadata.iloc[idx]
+        image_path = Path(self.images_root_path) / row.image_path
+        img = Image.open(image_path).convert("RGB")
+        img = v2.ToTensor()(img)
+        return {"features": img, "observation_id": row.observation_id}
+
+
+class InferenceDataModel(pl.LightningDataModule):
+    def __init__(
+        self,
+        metadata_path,
+        images_root_path,
+        batch_size=32,
+    ):
+        super().__init__()
+        self.metadata_path = metadata_path
+        self.images_root_path = images_root_path
+        self.batch_size = batch_size
+
+    def setup(self, stage=None):
+        self.dataloader = DataLoader(
+            ImageDataset(self.metadata_path, self.images_root_path),
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=4,
+        )
+
+    def predict_dataloader(self):
+        transform = v2.Compose([TransformDino("facebook/dinov2-base")])
+        for batch in self.dataloader:
+            batch = transform(batch)
+            yield batch

model.py

@@ -0,0 +1,81 @@
+import pytorch_lightning as pl
+import torch
+from torch import nn
+from torchmetrics.classification import (
+    MulticlassAccuracy,
+    MulticlassF1Score,
+    MulticlassPrecision,
+    MulticlassRecall,
+)
+
+
+class LinearClassifier(pl.LightningModule):
+    def __init__(self, num_features, num_classes):
+        super().__init__()
+        self.num_features = num_features
+        self.num_classes = num_classes
+        self.save_hyperparameters()  # Saves hyperparams in the checkpoints
+        self.model = nn.Linear(num_features, num_classes)
+        self.learning_rate = 0.002
+        self.accuracy = MulticlassAccuracy(num_classes=num_classes, average="weighted")
+        self.f1_score = MulticlassF1Score(num_classes=num_classes, average="weighted")
+        self.precision = MulticlassPrecision(
+            num_classes=num_classes, average="weighted"
+        )
+        self.recall = MulticlassRecall(num_classes=num_classes, average="weighted")
+
+    def forward(self, x):
+        return torch.log_softmax(self.model(x), dim=1)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        return optimizer
+
+    def _run_step(self, batch, batch_idx, step_name):
+        x, y = batch["features"], batch["label"]
+        logits = self(x)
+        loss = torch.nn.functional.nll_loss(logits, y)
+        self.log(f"{step_name}_loss", loss, prog_bar=True)
+        self.log(
+            f"{step_name}_accuracy",
+            self.accuracy(logits, y),
+            on_step=False,
+            on_epoch=True,
+        )
+        if step_name != "train":
+            self.log(
+                f"{step_name}_f1",
+                self.f1_score(logits, y),
+                on_step=False,
+                on_epoch=True,
+            )
+            self.log(
+                f"{step_name}_precision",
+                self.precision(logits, y),
+                on_step=False,
+                on_epoch=True,
+            )
+            self.log(
+                f"{step_name}_recall",
+                self.recall(logits, y),
+                on_step=False,
+                on_epoch=True,
+            )
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        return self._run_step(batch, batch_idx, "train")
+
+    def validation_step(self, batch, batch_idx):
+        return self._run_step(batch, batch_idx, "val")
+
+    def test_step(self, batch, batch_idx):
+        return self._run_step(batch, batch_idx, "test")
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=None):
+        logits = self(batch["features"])
+        return {
+            "logits": logits,
+            "class_id": torch.argmax(logits, dim=1),
+            "observation_id": batch["observation_id"],
+        }

submission.py

@@ -0,0 +1,41 @@
+import zipfile
+
+import pandas as pd
+import torch
+from pytorch_lightning import Trainer
+
+from .data import InferenceDataModel
+from .model import LinearClassifier
+
+
+def make_submission(
+    test_metadata,
+    model_path,
+    output_csv_path="./submission.csv",
+    images_root_path="/tmp/data/private_testset",
+):
+    model = LinearClassifier.load_from_checkpoint(model_path)
+    dm = InferenceDataModel(
+        metadata_path=test_metadata, images_root_path=images_root_path
+    )
+    trainer = Trainer(
+        accelerator="gpu" if torch.cuda.is_available() else "cpu",
+    )
+    predictions = trainer.predict(model, datamodule=dm)
+    rows = []
+    for batch in predictions:
+        for observation_id, class_id in zip(batch["observation_id"], batch["class_id"]):
+            row = {"observation_id": int(observation_id), "class_id": int(class_id)}
+            rows.append(row)
+    submission_df = pd.DataFrame(rows)
+    submission_df.to_csv(output_csv_path, index=False)
+
+
+if __name__ == "__main__":
+    with zipfile.ZipFile("/tmp/data/private_testset.zip", "r") as zip_ref:
+        zip_ref.extractall("/tmp/data")
+
+    MODEL_PATH = "last.ckpt"
+    metadata_file_path = "./SnakeCLEF2024-TestMetadata.csv"
+    test_metadata = pd.read_csv(metadata_file_path)
+    make_submission(test_metadata=test_metadata, model_path=MODEL_PATH)

test_evaluate.py

@@ -0,0 +1,82 @@
+import numpy as np
+import pandas as pd
+import PIL
+import pytest
+import torch
+from pytorch_lightning import Trainer
+
+from .data import ImageDataset, InferenceDataModel
+from .model import LinearClassifier
+from .submission import make_submission
+
+
+class TestingInferenceDataModel(InferenceDataModel):
+    def train_dataloader(self):
+        for batch in self.predict_dataloader():
+            # add a label to the batch with classes from 0 to 9
+            batch["label"] = torch.randint(0, 10, (batch["features"].shape[0],))
+            yield batch
+
+
+@pytest.fixture
+def images_root(tmp_path):
+    images_root = tmp_path / "images"
+    images_root.mkdir()
+    for i in range(10):
+        img = PIL.Image.fromarray(
+            np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+        )
+        img.save(images_root / f"{i}.jpg")
+    return images_root
+
+
+@pytest.fixture
+def metadata(tmp_path, images_root):
+    res = []
+    for i, img in enumerate(images_root.glob("*.jpg")):
+        res.append({"image_path": img.name, "observation_id": i})
+    df = pd.DataFrame(res)
+    df.to_csv(tmp_path / "metadata.csv", index=False)
+    return tmp_path / "metadata.csv"
+
+
+@pytest.fixture
+def model_checkpoint(tmp_path, metadata, images_root):
+    model_checkpoint = tmp_path / "model.ckpt"
+    model = LinearClassifier(768, 10)
+    trainer = Trainer(max_epochs=1, fast_dev_run=True)
+    dm = TestingInferenceDataModel(metadata, images_root)
+    trainer.fit(model, dm)
+    trainer.save_checkpoint(model_checkpoint)
+    return model_checkpoint
+
+
+def test_image_dataset(images_root, metadata):
+    dataset = ImageDataset(metadata, images_root)
+    assert len(dataset) == 10
+    for i in range(10):
+        assert dataset[i]["features"].shape == torch.Size([3, 100, 100])
+
+
+def test_inference_datamodel(images_root, metadata):
+    batch_size = 5
+    model = InferenceDataModel(metadata, images_root, batch_size=batch_size)
+    model.setup()
+    assert len(model.dataloader) == 2
+    for batch in model.predict_dataloader():
+        assert set(batch.keys()) == {"features", "observation_id"}
+        assert batch["features"].shape == torch.Size([batch_size, 768])
+
+
+def test_model_checkpoint(model_checkpoint):
+    model = LinearClassifier.load_from_checkpoint(model_checkpoint)
+    assert model
+
+
+def test_make_submission(model_checkpoint, metadata, images_root, tmp_path):
+    output_csv_path = tmp_path / "submission.csv"
+    make_submission(metadata, model_checkpoint, output_csv_path, images_root)
+    submission_df = pd.read_csv(output_csv_pathgit)
+    assert len(submission_df) == 10
+    assert set(submission_df.columns) == {"observation_id", "class_id"}
+    assert submission_df["class_id"].isin(range(10)).all()