Update script.py

script.py

@@ -1,31 +1,45 @@
 import pandas as pd
 import numpy as np
-import onnxruntime as ort
 import os
+from torch import nn
 from tqdm import tqdm
 import timm
 import torchvision.transforms as T
 from PIL import Image
 import torch
 
+
 def is_gpu_available():
     """Check if the python package `onnxruntime-gpu` is installed."""
     return torch.cuda.is_available()
 
+class CustomModel(nn.Module):
+    def __init__(self, base_model_name, num_classes1, num_classes2):
+        super(CustomModel, self).__init__()
+        self.base_model = timm.create_model(base_model_name, pretrained=False)
+        in_features = self.base_model.get_classifier().in_features
+        self.base_model.reset_classifier(0)  # Remove the original classification layer
+
+        self.fc1 = nn.Linear(in_features, num_classes1)  # Binary classification output
+        self.fc2 = nn.Linear(in_features, num_classes2)  # Categorical classification output
+
+    def forward(self, x):
+        x = self.base_model(x)
+        out1 = torch.sigmoid(self.fc1(x))  # Binary output
+        out2 = self.fc2(x)  # Categorical output
+        return out2
 
 class PytorchWorker:
     """Run inference using ONNX runtime."""
 
     def __init__(self, model_path: str, model_name: str, number_of_categories: int = 1784):
-
         def _load_model(model_name, model_path):
-
             print("Setting up Pytorch Model")
             self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
             print(f"Using devide: {self.device}")
 
-            model = timm.create_model(model_name, num_classes=number_of_categories, pretrained=False)
-
+            # model = timm.create_model(model_name, num_classes=number_of_categories, pretrained=False)
+            model =CustomModel(model_name, 1, number_of_categories)
             # if not torch.cuda.is_available():
             #     model_ckpt = torch.load(model_path, map_location=torch.device("cpu"))
             # else:
@@ -42,10 +56,8 @@ class PytorchWorker:
                                      T.ToTensor(),
                                      T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
 
-
     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.
         """
@@ -55,7 +67,8 @@ 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_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)
@@ -78,9 +91,8 @@ def make_submission(test_metadata, model_path, model_name, output_csv_path="./su
 
 
 if __name__ == "__main__":
-
     import zipfile
-    
+
     with zipfile.ZipFile("/tmp/data/private_testset.zip", 'r') as zip_ref:
         zip_ref.extractall("/tmp/data")
 
@@ -99,3 +111,4 @@ if __name__ == "__main__":
         model_name=MODEL_NAME,
         # images_root_path='/home/zeleznyt/mnt/data-ntis/projects/korpusy_cv/SnakeCLEF2024/val/SnakeCLEF2023-medium_size'
     )
+