fix

script.py

@@ -121,7 +121,7 @@ class ImageDataset(Dataset):
         image = cv2.imread(image_path)
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
-        return self.transform(image=image)['image']
+        return self.transform(image=image)["image"]
 
 
 class EmbeddingMetadataDataset(Dataset):
@@ -152,16 +152,20 @@ class EmbeddingMetadataDataset(Dataset):
 
 def generate_embeddings(metadata_file_path, root_dir):
 
+    DINOV2_CKPT = "./checkpoints/dinov2.bin"
+
     metadata_df = pd.read_csv(metadata_file_path)
 
     test_dataset = ImageDataset(metadata_df, local_filepath=root_dir)
 
-    loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)
+    loader = DataLoader(test_dataset, batch_size=2, 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
+        "timm/vit_large_patch14_reg4_dinov2.lvd142m", pretrained=False
     )
+    weights = torch.load(DINOV2_CKPT)
+    model.load_state_dict(weights)
 
     model = model.to(device)
     model.eval()
@@ -255,9 +259,7 @@ class FungiMEEModel(nn.Module):
         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
-        ) 
+        full_emb = torch.stack((img_emb, date_emb, geo_emb, substr_emb), dim=1)
 
         cls_emb = self.encoder.forward(full_emb)[:, 0, :].squeeze(1)
 
@@ -305,37 +307,46 @@ class FungiEnsembleModel(nn.Module):
 
         return logits.argmax(1).tolist()
 
+
 def make_submission(metadata_df):
 
     OUTPUT_CSV_PATH = "./submission.csv"
     BASE_CKPT_PATH = "./checkpoints"
 
-    ckpt_path = os.path.join(BASE_CKPT_PATH, "dino_2_optuna_05242055.ckpt")
-    # model_names = os.listdir(BASE_CKPT_PATH)
-
-    # models = []
-
-    # for model_path in model_names:
-    #     print("loading ", model_path)
-    #     ckpt_path = os.path.join(BASE_CKPT_PATH, model_path)
-
-    #     ckpt = torch.load(ckpt_path)
-    #     model = FungiMEEModel()
-    #     model.load_state_dict(
-    #         {w: ckpt["model." + w] for w in model.state_dict().keys()}
-    #     )
-    #     model.eval()
-    #     model.cuda()
+    # ckpt_path = os.path.join(BASE_CKPT_PATH, "dino_2_optuna_05242055.ckpt")
+    model_names = [
+        "dino_2_optuna_05242231.ckpt",
+        "dino_optuna_05241449.ckpt",
+        "dino_optuna_05241257.ckpt",
+        "dino_optuna_05241222.ckpt",
+        "dino_2_optuna_05242055.ckpt",
+        "dino_2_optuna_05242156.ckpt",
+        "dino_2_optuna_05242344.ckpt",
+    ]
+
+    models = []
+
+    for model_path in model_names:
+        print("loading ", model_path)
+        ckpt_path = os.path.join(BASE_CKPT_PATH, model_path)
+
+        ckpt = torch.load(ckpt_path)
+        model = FungiMEEModel()
+        model.load_state_dict(
+            {w: ckpt["model." + w] for w in model.state_dict().keys()}
+        )
+        model.eval()
+        model.cuda()
 
-    #     models.append(model)
+        models.append(model)
 
-    # fungi_model = FungiEnsembleModel(models)
+    fungi_model = FungiEnsembleModel(models)
 
-    fungi_model = FungiMEEModel()
-    ckpt = torch.load(ckpt_path)
-    fungi_model.load_state_dict(
-        {w: ckpt["model." + w] for w in fungi_model.state_dict().keys()}
-    )
+    # fungi_model = FungiMEEModel()
+    # ckpt = torch.load(ckpt_path)
+    # fungi_model.load_state_dict(
+    #     {w: ckpt["model." + w] for w in fungi_model.state_dict().keys()}
+    # )
 
     embedding_dataset = EmbeddingMetadataDataset(metadata_df)
     loader = DataLoader(embedding_dataset, batch_size=128, shuffle=False)
@@ -369,18 +380,18 @@ if __name__ == "__main__":
     MODEL_PATH = "metaformer-s-224.pth"
     MODEL_NAME = "timm/vit_base_patch14_reg4_dinov2.lvd142m"
 
-    # # # # Real submission
-    # import zipfile
+    # # # # # Real submission
+    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"
 
     ##############