fixed script

script.py

@@ -1,4 +1,3 @@
-import io
 import os
 from typing import List
 
@@ -10,9 +9,18 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.transforms as T
-from albumentations import (CenterCrop, Compose, HorizontalFlip, Normalize,
-                            PadIfNeeded, RandomBrightnessContrast, RandomCrop,
-                            RandomResizedCrop, Resize, VerticalFlip)
+from albumentations import (
+    CenterCrop,
+    Compose,
+    HorizontalFlip,
+    Normalize,
+    PadIfNeeded,
+    RandomBrightnessContrast,
+    RandomCrop,
+    RandomResizedCrop,
+    Resize,
+    VerticalFlip,
+)
 from albumentations.pytorch import ToTensorV2
 from PIL import Image
 from timm.layers import LayerNorm2d, SelectAdaptivePool2d
@@ -20,14 +28,13 @@ from timm.models.metaformer import MlpHead
 from torch.utils.data import DataLoader, Dataset
 from tqdm import tqdm
 
-DEFAULT_WIDTH = 518
-DEFAULT_HEIGHT = 518
+DIM = 518
 
 def get_transforms(*, data, model=None, width=None, height=None):
     assert data in ("train", "valid")
 
-    width = width if width else DEFAULT_WIDTH
-    height = height if height else DEFAULT_HEIGHT
+    width = width if width else DIM
+    height = height if height else DIM
 
     model_mean = list(model.default_cfg["mean"]) if model else (0.5, 0.5, 0.5)
     model_std = list(model.default_cfg["std"]) if model else (0.5, 0.5, 0.5)
@@ -53,8 +60,6 @@ def get_transforms(*, data, model=None, width=None, height=None):
             ]
         )
 
-DIM = 518
-BASE_PATH = "../data/DF_FULL"
 
 def generate_embeddings(metadata_file_path, root_dir):
 
@@ -69,7 +74,9 @@ def generate_embeddings(metadata_file_path, root_dir):
     loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model = timm.create_model("timm/vit_large_patch14_reg4_dinov2.lvd142m", pretrained=True)
+    model = timm.create_model(
+        "timm/vit_large_patch14_reg4_dinov2.lvd142m", pretrained=True
+    )
     model = model.to(device)
     model.eval()
 
@@ -87,13 +94,14 @@ def generate_embeddings(metadata_file_path, root_dir):
 
     embs_list = [x for x in all_embs]
     metadata_df["embedding"] = embs_list
-    
+
     return metadata_df
 
 
-TIME = ['m0', 'm1', 'd0', 'd1']
-GEO = ['g0', 'g1', 'g2', 'g3', 'g4', 'g5', 'g_float']
-SUBSTRATE = ["substrate_0",
+TIME = ["m0", "m1", "d0", "d1"]
+GEO = ["g0", "g1", "g2", "g3", "g4", "g5", "g_float"]
+SUBSTRATE = [
+    "substrate_0",
     "substrate_1",
     "substrate_2",
     "substrate_3",
@@ -168,11 +176,12 @@ SUBSTRATE = ["substrate_0",
     "habitat_31",
 ]
 
+
 class EmbeddingMetadataDataset(Dataset):
     def __init__(self, df):
         self.df = df
 
-        self.emb = df['embedding']
+        self.emb = df["embedding"]
         self.metadata_date = df[TIME].to_numpy()
         self.metadata_geo = df[GEO].to_numpy()
         self.metadata_substrate = df[SUBSTRATE].to_numpy()
@@ -186,24 +195,27 @@ class EmbeddingMetadataDataset(Dataset):
         metadata = {
             "date": torch.from_numpy(self.metadata_date[idx, :]).type(torch.float),
             "geo": torch.from_numpy(self.metadata_geo[idx, :]).type(torch.float),
-            "substr": torch.from_numpy(self.metadata_substrate[idx, :]).type(torch.float),
+            "substr": torch.from_numpy(self.metadata_substrate[idx, :]).type(
+                torch.float
+            ),
         }
 
         return embedding, metadata
-    
+
 
 class ImageMetadataDataset(Dataset):
     def __init__(self, df, transform=None, local_filepath=None):
         self.df = df
         self.transform = transform
         self.local_filepath = local_filepath
-        
-        self.filepaths = df["image_path"].apply(lambda x: x.replace("jpg", "JPG")).to_list()
+
+        self.filepaths = (
+            df["image_path"].to_list()
+        )
         self.metadata_date = df[TIME].to_numpy()
         self.metadata_geo = df[GEO].to_numpy()
         self.metadata_substrate = df[SUBSTRATE].to_numpy()
 
-
     def __len__(self):
         return len(self.df)
 
@@ -223,16 +235,20 @@ class ImageMetadataDataset(Dataset):
         metadata = {
             "date": torch.from_numpy(self.metadata_date[idx, :]).type(torch.float),
             "geo": torch.from_numpy(self.metadata_geo[idx, :]).type(torch.float),
-            "substr": torch.from_numpy(self.metadata_substrate[idx, :]).type(torch.float),
+            "substr": torch.from_numpy(self.metadata_substrate[idx, :]).type(
+                torch.float
+            ),
         }
 
         return image, metadata
 
+
 DATE_SIZE = 4
 GEO_SIZE = 7
 SUBSTRATE_SIZE = 73
 NUM_CLASSES = 1717
 
+
 class StarReLU(nn.Module):
     """
     StarReLU: s * relu(x) ** 2 + b
@@ -260,6 +276,7 @@ class StarReLU(nn.Module):
     def forward(self, x):
         return self.scale * self.relu(x) ** 2 + self.bias
 
+
 class FungiMEEModel(nn.Module):
     def __init__(
         self,
@@ -272,7 +289,6 @@ class FungiMEEModel(nn.Module):
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         print(f"Using devide: {self.device}")
 
-
         self.date_embedding = MlpHead(
             dim=DATE_SIZE, num_classes=dim, mlp_ratio=128, act_layer=StarReLU
         )
@@ -286,38 +302,43 @@ class FungiMEEModel(nn.Module):
             act_layer=StarReLU,
         )
 
-        self.encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=dim, nhead=8, batch_first=True), num_layers=4)
-        
+        self.encoder = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(d_model=dim, nhead=8, batch_first=True),
+            num_layers=4,
+        )
+
         self.head = MlpHead(dim=dim, num_classes=num_classes, drop_rate=0)
 
         for param in self.parameters():
             if param.dim() > 1:
                 nn.init.kaiming_normal_(param)
 
-
     def forward(self, img_emb, metadata):
 
         img_emb = img_emb.to(self.device)
-        
+
         date_emb = self.date_embedding.forward(metadata["date"].to(self.device))
         geo_emb = self.geo_embedding.forward(metadata["geo"].to(self.device))
         substr_emb = self.substr_embedding.forward(metadata["substr"].to(self.device))
 
-        full_emb = torch.stack((img_emb, date_emb, geo_emb, substr_emb), dim=1) #.unsqueeze(0)
+        full_emb = torch.stack(
+            (img_emb, date_emb, geo_emb, substr_emb), dim=1
+        )  # .unsqueeze(0)
         # print(full_emb.shape)
 
         cls_emb = self.encoder.forward(full_emb)[:, 0, :].squeeze(1)
 
         return self.head.forward(cls_emb)
-    
+
     def predict(self, img_emb, metadata):
-        
+
         logits = self.forward(img_emb, metadata)
 
         # Any preprocess happens here
 
         return logits.argmax(1).tolist()
-    
+
+
 class FungiEnsembleModel(nn.Module):
 
     def __init__(self, models, softmax=True) -> None:
@@ -331,39 +352,46 @@ class FungiEnsembleModel(nn.Module):
             model = model.to(self.device)
             model.eval()
             self.models.append(model)
-        
+
     def forward(self, img_emb, metadata):
 
         img_emb = img_emb.to(self.device)
 
-        probs = []        
+        probs = []
 
         for model in self.models:
             logits = model.forward(img_emb, metadata)
-            
-            p = logits.softmax(dim=1).detach().cpu() if self.softmax else logits.detach().cpu()
+
+            p = (
+                logits.softmax(dim=1).detach().cpu()
+                if self.softmax
+                else logits.detach().cpu()
+            )
 
             probs.append(p)
 
         return torch.stack(probs).mean(dim=0)
-    
+
     def predict(self, img_emb, metadata):
-        
+
         logits = self.forward(img_emb, metadata)
 
         # Any preprocess happens here
 
         return logits.argmax(1).tolist()
-    
+
 
 def is_gpu_available():
     """Check if the python package `onnxruntime-gpu` is installed."""
     return torch.cuda.is_available()
 
+
 class PytorchWorker:
     """Run inference using ONNX runtime."""
 
-    def __init__(self, model_path: str, model_name: str, number_of_categories: int = 1605):
+    def __init__(
+        self, model_path: str, model_name: str, number_of_categories: int = 1605
+    ):
 
         def _load_model(model_name, model_path):
 
@@ -379,10 +407,13 @@ class PytorchWorker:
 
         self.model = _load_model(model_name, model_path)
 
-        self.transforms = T.Compose([T.Resize((518, 518)),
-                                     T.ToTensor(),
-                                     T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
-
+        self.transforms = T.Compose(
+            [
+                T.Resize((518, 518)),
+                T.ToTensor(),
+                T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+            ]
+        )
 
     def predict_image(self, image: np.ndarray):
         """Run inference using ONNX runtime.
@@ -397,9 +428,9 @@ class PytorchWorker:
 
 
 def make_submission(metadata_df, model_names=None):
-    
-    OUTPUT_CSV_PATH="./submission.csv"
-    
+
+    OUTPUT_CSV_PATH = "./submission.csv"
+
     """Make submission with given """
 
     BASE_CKPT_PATH = "./checkpoints"
@@ -414,12 +445,14 @@ def make_submission(metadata_df, model_names=None):
 
         ckpt = torch.load(ckpt_path)
         model = FungiMEEModel()
-        model.load_state_dict({w: ckpt['state_dict']["model." + w] for w in model.state_dict().keys()})
+        model.load_state_dict(
+            {w: ckpt["state_dict"]["model." + w] for w in model.state_dict().keys()}
+        )
         model.eval()
         model.cuda()
 
         models.append(model)
-    
+
     ensemble_model = FungiEnsembleModel(models)
 
     embedding_dataset = EmbeddingMetadataDataset(metadata_df)
@@ -433,31 +466,39 @@ def make_submission(metadata_df, model_names=None):
 
     all_preds = torch.vstack(preds).numpy()
 
-    preds_df = metadata_df[['observation_id', 'image_path']]
-    preds_df['preds'] = [i for i in all_preds]
-    preds_df = preds_df[['observation_id', 'preds']].groupby('observation_id').mean().reset_index()
-    preds_df['class_id'] = preds_df['preds'].apply(lambda x: x.argmax() if x.argmax() <= 1603 else -1)
-    preds_df[['observation_id', 'class_id']].to_csv(OUTPUT_CSV_PATH, index=None)
+    preds_df = metadata_df[["observation_id", "image_path"]]
+    preds_df["preds"] = [i for i in all_preds]
+    preds_df = (
+        preds_df[["observation_id", "preds"]]
+        .groupby("observation_id")
+        .mean()
+        .reset_index()
+    )
+    preds_df["class_id"] = preds_df["preds"].apply(
+        lambda x: x.argmax() if x.argmax() <= 1603 else -1
+    )
+    preds_df[["observation_id", "class_id"]].to_csv(OUTPUT_CSV_PATH, index=None)
 
     print("Submission complete")
 
+
 if __name__ == "__main__":
 
     MODEL_PATH = "metaformer-s-224.pth"
     MODEL_NAME = "timm/vit_base_patch14_reg4_dinov2.lvd142m"
 
     # # Real submission
-    import zipfile
+    # import zipfile
 
-    with zipfile.ZipFile("/tmp/data/private_testset.zip", 'r') as zip_ref:
-        zip_ref.extractall("/tmp/data")
+    # with zipfile.ZipFile("/tmp/data/private_testset.zip", "r") as zip_ref:
+    #     zip_ref.extractall("/tmp/data")
 
-    metadata_file_path = "./_test_preprocessed.csv"
-    root_dir = "/tmp/data"
+    # metadata_file_path = "./_test_preprocessed.csv"
+    # root_dir = "/tmp/data"
 
     # Test submission
-    # metadata_file_path = "../trial_submission.csv"
-    # root_dir = "../data/DF_FULL"
+    metadata_file_path = "../trial_submission.csv"
+    root_dir = "../data/DF_FULL"
 
     ##############