添加模型

frame_field_learning/inference.py

@@ -5,7 +5,7 @@ import scipy
 
 import numpy as np
 import torch
-
+from huggingface_hub import hf_hub_download
 from . import local_utils
 from . import polygonize
 
@@ -145,20 +145,20 @@ def load_checkpoint(model, checkpoints_dirpath, device):
     """
     Loads best val checkpoint in checkpoints_dirpath
     """
-    filepaths = python_utils.get_filepaths(checkpoints_dirpath, startswith_str="checkpoint.best_val.",
-                                           endswith_str=".tar")
-    if len(filepaths):
-        filepaths = sorted(filepaths)
-        filepath = filepaths[-1]  # Last best val checkpoint filepath in case there is more than one
-        print_utils.print_info("Loading best val checkpoint: {}".format(filepath))
-    else:
-        # No best val checkpoint fount: find last checkpoint:
-        filepaths = python_utils.get_filepaths(checkpoints_dirpath, endswith_str=".tar",
-                                               startswith_str="checkpoint.")
-        filepaths = sorted(filepaths)
-        filepath = filepaths[-1]  # Last checkpoint
-        print_utils.print_info("Loading last checkpoint: {}".format(filepath))
-
+    # filepaths = python_utils.get_filepaths(checkpoints_dirpath, startswith_str="checkpoint.best_val.",
+    #                                        endswith_str=".tar")
+    # if len(filepaths):
+    #     filepaths = sorted(filepaths)
+    #     filepath = filepaths[-1]  # Last best val checkpoint filepath in case there is more than one
+    #     print_utils.print_info("Loading best val checkpoint: {}".format(filepath))
+    # else:
+    #     # No best val checkpoint fount: find last checkpoint:
+    #     filepaths = python_utils.get_filepaths(checkpoints_dirpath, endswith_str=".tar",
+    #                                            startswith_str="checkpoint.")
+    #     filepaths = sorted(filepaths)
+    #     filepath = filepaths[-1]  # Last checkpoint
+    #     print_utils.print_info("Loading last checkpoint: {}".format(filepath))
+    filepath = hf_hub_download(repo_id="Egrt/Luuuu", filename="checkpoint.best_val.epoch_000047.tar")
     device = torch.device(device)
     checkpoint = torch.load(filepath, map_location=device)  # map_location is used to load on current device
 

requirements.txt

@@ -1,4 +1,5 @@
 cython
+huggingface_hub
 scipy==1.4.1
 numpy==1.22.3
 matplotlib==3.3.2