Update script.py

script.py

@@ -55,10 +55,11 @@ class PytorchWorker:
         return logits.tolist()
 
 
-def make_submission(test_metadata, model_path, model_name, output_csv_path="./submission.csv", images_root_path="/tmp/data/private_testset"):
+def make_submission(test_metadata, model_paths, model_name, output_csv_path="./submission.csv", images_root_path="/tmp/data/private_testset"):
     """Make submission with given """
-
-    model = PytorchWorker(model_path, model_name)
+    models = []
+    for m in model_paths:
+        models.append(PytorchWorker(m, model_name))
 
     predictions = []
 
@@ -66,8 +67,15 @@ def make_submission(test_metadata, model_path, model_name, output_csv_path="./su
         image_path = os.path.join(images_root_path, row.filename)
 
         test_image = Image.open(image_path).convert("RGB")
+        flipped_image = test_image.transpose(Image.FLIP_LEFT_RIGHT)
+
+        result_logits = []
+
+        for model in models:
+            result_logits += model.predict_image(test_image)
+            # result_logits += model.predict_image(flipped_image)
 
-        logits = model.predict_image(test_image)
+        logits = np.average(np.array(result_logits), 0)
 
         predictions.append(np.argmax(logits))
 
@@ -85,7 +93,11 @@ if __name__ == "__main__":
         zip_ref.extractall("/tmp/data")
 
     # MODEL_PATH = "pytorch_model.bin"
-    MODEL_PATH = "2405_PSCmetric_real3_best_accuracy.pth"
+    MODEL_PATH = ["2405_PSCmetric2_0001_best_accuracy.pth",
+                  "2405_PSCmetric_3x_best_accuracy.pth",
+                  "2405_PSCmetric_real3_best_accuracy.pth",
+                  "2405_cls_boost_best_accuracy.pth"
+                 ]
     # MODEL_NAME = "tf_efficientnet_b1.ap_in1k"
     MODEL_NAME = "swinv2_tiny_window16_256.ms_in1k"
 
@@ -95,8 +107,7 @@ if __name__ == "__main__":
 
     make_submission(
         test_metadata=test_metadata,
-        model_path=MODEL_PATH,
+        model_paths=MODEL_PATH,
         model_name=MODEL_NAME,
         # images_root_path='/home/zeleznyt/mnt/data-ntis/projects/korpusy_cv/SnakeCLEF2024/val/SnakeCLEF2023-medium_size'
     )
-s