Update script.py

script.py

@@ -59,16 +59,15 @@ class PytorchWorker:
         return outputs.cpu()  # Convert tensor to list
 
 
-def make_submission(test_metadata, model_path, model_path2, model_name, output_csv_path="./submission.csv", images_root_path="/tmp/data/private_testset"):
+def make_submission(test_metadata, model_path, model_path2, model_name, model_name2, output_csv_path="./submission.csv", images_root_path="/tmp/data/private_testset"):
     """Make submission with given """
-
-    model = PytorchWorker(model_path, model_name)
-    model2 = PytorchWorker(model_path2, model_name)
+    model = PytorchWorker(model_path, model_name, number_of_categories=1604)
+    model2 = PytorchWorker(model_path2, model_name2)
     
     predictions = []
     correct_max_values = []
     incorrect_max_values = []
-
+    
     for _, row in tqdm(test_metadata.iterrows(), total=len(test_metadata)):
         image_path = os.path.join(images_root_path, row.image_path)
         test_image = Image.open(image_path).convert("RGB")
@@ -77,11 +76,14 @@ def make_submission(test_metadata, model_path, model_path2, model_name, output_c
         outputs2 = model2.predict_image(test_image)
         
         # max_value = torch.max(outputs+outputs2)
-        _, preds = torch.max(outputs+outputs2, 1)
+        _, preds = torch.max(outputs, 1) # baseline
+        _, preds2 = torch.max(outputs2, 1) # 1.4.3
+        
         pred_class_id = preds.tolist()
+        pred_class_id2 = preds2.tolist()
         # max_value2 = torch.max(outputs2)
         
-        pred_class_id = pred_class_id[0] if pred_class_id[0] != 1604 else -1
+        pred_class_id = pred_class_id[0] if pred_class_id2[0] != 1604 else -1
         
         predictions.append(pred_class_id)
 
@@ -105,10 +107,10 @@ if __name__ == "__main__":
     # MODEL_PATH = './efficientnet_b3_epoch_21_trick1.2.5_a0.7237_l17.1662.pth'
     # MODEL_PATH = './efficientnet_b3_epoch_21_trcik1.5.2.pth'
     # MODEL_PATH = './efficientnet_b3_epoch_28_1.4.3.pth'
-    MODEL_PATH = './efficientnet_b3_epoch_28_1.4.3.pth'
-    MODEL_PATH2 = './efficientnet_b3_epoch_23_trick1.4.1.pth'
-    MODEL_NAME = 'tf_efficientnet_b3_ns' #"tf_efficientnet_b1.ap_in1k"
-
+    MODEL_PATH = './pytorch_model.bin'
+    MODEL_PATH2 = './efficientnet_b3_epoch_28_1.4.3.pth'
+    MODEL_NAME = "tf_efficientnet_b1_ap"
+    MODEL_NAME2 = "tf_efficientnet_b3_ns"
     metadata_file_path = "./FungiCLEF2024_TestMetadata.csv"
     test_metadata = pd.read_csv(metadata_file_path)
 
@@ -116,5 +118,6 @@ if __name__ == "__main__":
         test_metadata=test_metadata,
         model_path=MODEL_PATH,
         model_path2=MODEL_PATH2,
-        model_name=MODEL_NAME
-    )
+        model_name=MODEL_NAME,
+        model_name2=MODEL_NAME2
+    )