Update script.py

script.py

@@ -49,41 +49,47 @@ class PytorchWorker:
                                     #  T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
 
 
-    def predict_image(self, image: np.ndarray) -> list:
+    def predict_image(self, image: np.ndarray) -> list():
         """Run inference using ONNX runtime.
         :param image: Input image as numpy array.
         :return: A list with logits and confidences.
         """
+
+        # logits = self.model(self.transforms(image).unsqueeze(0).to(self.device))
+        
         self.model.eval()
+        
         outputs = self.model(self.transforms(image).unsqueeze(0).to(self.device))
-        return outputs.cpu()  # Convert tensor to list
+        
+        _, preds = torch.max(outputs, 1)
+        
+        preds = preds.cpu()  # Move tensor to CPU    
 
+        # post process
+        # max_value = torch.max(outputs)
+        # if max_value < -20:
+        #     preds[0]=1604
+        
+        print("preds: ", preds)
+
+        return preds.tolist()  # Convert tensor to list
 
-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"):
+
+def make_submission(test_metadata, model_path, model_name, output_csv_path="./submission.csv", images_root_path="/tmp/data/private_testset"):
     """Make submission with given """
-    model = PytorchWorker(model_path, model_name, number_of_categories=1604)
-    model2 = PytorchWorker(model_path2, model_name2)
-    
+
+    model = PytorchWorker(model_path, model_name)
+
     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")
+
+        logits = model.predict_image(test_image)
         
-        outputs = model.predict_image(test_image)
-        outputs2 = model2.predict_image(test_image)
-        
-        # max_value = torch.max(outputs+outputs2)
-        _, 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_id2[0] != 1604 else -1
+        pred_class_id = logits[0] if logits[0] !=1604 else -1
         
         predictions.append(pred_class_id)
 
@@ -107,17 +113,15 @@ 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 = './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"
+    # MODEL_PATH = './efficientnet_b3_epoch_28_trick1.4.3.2.pth'
+    MODEL_PATH = './fused_model_soup.pth'
+    MODEL_NAME = 'tf_efficientnet_b3_ns' #"tf_efficientnet_b1.ap_in1k"
+
     metadata_file_path = "./FungiCLEF2024_TestMetadata.csv"
     test_metadata = pd.read_csv(metadata_file_path)
 
     make_submission(
         test_metadata=test_metadata,
         model_path=MODEL_PATH,
-        model_path2=MODEL_PATH2,
-        model_name=MODEL_NAME,
-        model_name2=MODEL_NAME2
+        model_name=MODEL_NAME
     )