Move everything into a single script

evaluate/submission.py

@@ -1,10 +1,31 @@
+from pathlib import Path
+
 import pandas as pd
 import torch
+from PIL import Image
 from torch import nn
-from torch.utils.data import DataLoader
-from torchvision.transforms import v2
+from torch.utils.data import DataLoader, Dataset
+from transformers import AutoImageProcessor, AutoModel
+import numpy as np
+
+
+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)
 
-from .data import ImageDataset, TransformDino
+    def __getitem__(self, idx):
+        row = self.metadata.iloc[idx]
+        image_path = Path(self.images_root_path) / row.filename
+        img = Image.open(image_path).convert("RGB")
+        # convert to numpy array
+        img = torch.from_numpy(np.array(img))
+        # img = torch.tensor(img).permute(2, 0, 1).float() / 255.0
+        return {"features": img, "observation_id": row.observation_id}
 
 
 class LinearClassifier(nn.Module):
@@ -18,6 +39,21 @@ class LinearClassifier(nn.Module):
         return torch.log_softmax(self.model(x), dim=1)
 
 
+class TransformDino:
+    def __init__(self, model_name="facebook/dinov2-base"):
+        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
+
+
 def make_submission(
     test_metadata,
     model_path,
@@ -29,13 +65,13 @@ def make_submission(
     model = LinearClassifier(hparams["num_features"], hparams["num_classes"])
     model.load_state_dict(checkpoint["state_dict"])
 
-    transform = v2.Compose([TransformDino("facebook/dinov2-base")])
+    transform = TransformDino()
     dataloader = DataLoader(
         ImageDataset(test_metadata, images_root_path), batch_size=32, num_workers=4
     )
     rows = []
     for batch in dataloader:
-        batch = transform(batch)
+        batch = transform.forward(batch)
         observation_ids = batch["observation_id"]
         logits = model(batch["features"])
         class_ids = torch.argmax(logits, dim=1)

script.py

@@ -1,7 +1,86 @@
 #!/usr/bin/env python
 import zipfile
-from evaluate.submission import make_submission
 from argparse import ArgumentParser
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from transformers import AutoImageProcessor, AutoModel
+
+
+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.filename
+        img = Image.open(image_path)
+        img = torch.from_numpy(np.array(img))
+        return {"features": img, "observation_id": row.observation_id}
+
+
+class LinearClassifier(nn.Module):
+    def __init__(self, num_features, num_classes):
+        super().__init__()
+        self.num_features = num_features
+        self.num_classes = num_classes
+        self.model = nn.Linear(num_features, num_classes)
+
+    def forward(self, x):
+        return torch.log_softmax(self.model(x), dim=1)
+
+
+class TransformDino:
+    def __init__(self, model_name="facebook/dinov2-base"):
+        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
+
+
+def make_submission(
+    test_metadata,
+    model_path,
+    output_csv_path="./submission.csv",
+    images_root_path="/tmp/data/private_testset",
+):
+    checkpoint = torch.load(model_path)
+    hparams = checkpoint["hyper_parameters"]
+    model = LinearClassifier(hparams["num_features"], hparams["num_classes"])
+    model.load_state_dict(checkpoint["state_dict"])
+
+    transform = TransformDino()
+    dataloader = DataLoader(
+        ImageDataset(test_metadata, images_root_path), batch_size=32, num_workers=4
+    )
+    rows = []
+    for batch in dataloader:
+        batch = transform.forward(batch)
+        observation_ids = batch["observation_id"]
+        logits = model(batch["features"])
+        class_ids = torch.argmax(logits, dim=1)
+        for observation_id, class_id in zip(observation_ids, class_ids):
+            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)
 
 
 def parse_args():