fix: fixing problem with HF hub

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 FungiCLEF2024_TestMetadata.csv filter=lfs diff=lfs merge=lfs -text
+pytorch_model.bin filter=lfs diff=lfs merge=lfs -text

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25dcfebee82c8b14a9a43b4934173becdc434745344257046a4f3eb4fe94dc6f
+size 34696073

script.py

@@ -16,20 +16,32 @@ def is_gpu_available():
 class PytorchWorker:
     """Run inference using ONNX runtime."""
 
-    def __init__(self, onnx_path: str):
-        print("Setting up Pytorch Model")
+    def __init__(self, model_path: str, model_name: str, number_of_categories: int = 1604):
 
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        def _load_model(model_name, model_path):
 
-        print(f"Using devide: {self.device}")
-        self.model = timm.create_model("hf-hub:BVRA/tf_efficientnet_b3.in1k_ft_df20_224", pretrained=True)
-        self.model = self.model.eval()
-        self.model.to(self.device)
+            print("Setting up Pytorch Model")
+            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+            print(f"Using devide: {device}")
+
+            model = timm.create_model(model_name, num_classes=number_of_categories, pretrained=False)
+
+            if not torch.cuda.is_available():
+                model_ckpt = torch.load(model_path, map_location=torch.device("cpu"))
+            else:
+                model_ckpt = torch.load(model_path)["model"]
+
+            model.load_state_dict(model_ckpt)
+
+            return model.to(device).eval()
+
+        self.model = _load_model(model_name, model_path)
 
         self.transforms = T.Compose([T.Resize((224, 224)),
                                      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.
 
@@ -42,10 +54,10 @@ class PytorchWorker:
         return logits.tolist()
 
 
-def make_submission(test_metadata, model_path, 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 = PytorchWorker(model_path, model_name)
 
     predictions = []
 
@@ -71,12 +83,14 @@ if __name__ == "__main__":
     with zipfile.ZipFile("/tmp/data/private_testset.zip", 'r') as zip_ref:
         zip_ref.extractall("/tmp/data")
 
-    HFHUB_MODEL_PATH = "hf-hub:BVRA/tf_efficientnet_b3.in1k_ft_df20_224"
+    MODEL_PATH = "pytorch_model.bin"
+    MODEL_NAME = "tf_efficientnet_b1.ap_in1k"
 
     metadata_file_path = "./FungiCLEF2024_TestMetadata.csv"
-    test_metadata = pd.read_csv(metadata_file_path)
+    test_metadata = pd.read_csv(metadata_file_path)[:100]
 
     make_submission(
         test_metadata=test_metadata,
-        model_path=HFHUB_MODEL_PATH,
+        model_path=MODEL_PATH,
+        model_name=MODEL_NAME
     )