Enhance Checkpoint Management, Debugging, and File Utilities

.circleci/data

@@ -0,0 +1 @@
+/home/ti_wang/data

deeplabcut/common_tools/download.py

@@ -0,0 +1,58 @@
+from huggingface_hub import hf_hub_download
+import os
+
+def download_from_huggingface_hub(target_folder_path, repo_id, filename, subfolder=None):
+    """
+    Download a file from the Hugging Face Hub to a specified local directory.
+
+    Parameters:
+        target_folder_path (str): Local directory path where file will be saved
+        repo_id (str): Hugging Face repository ID (e.g., 'noahcao/sapiens-pose-coco')
+        filename (str): Name of the file to download
+        subfolder (str, optional): Path to subfolder within the repository where the file is located
+
+    Returns:
+        str: Full path to the downloaded file
+
+    Examples:
+        >>> # Download a model file from noahcao/sapiens-pose-coco
+        >>> download_from_huggingface_hub(
+        ...     target_folder_path="./models/sapiens",
+        ...     repo_id="noahcao/sapiens-pose-coco",
+        ...     filename="sapiens_2b_coco_best_coco_AP_822_torchscript.pt2",
+        ...     subfolder="sapiens_lite_host/torchscript/pose/checkpoints/sapiens_2b"
+        ... )
+
+        >>> # Download a file without specifying a subfolder
+        >>> download_from_huggingface_hub(
+        ...     target_folder_path="./data",
+        ...     repo_id="noahcao/sapiens-pose-coco",
+        ...     filename="COCO_val2017_detections_AP_H_70_person.json",
+        ...     subfolder="sapiens_host/pose/person_detection_results"
+        ... )
+    """
+    # Create the target directory if it does not exist
+    os.makedirs(target_folder_path, exist_ok=True)
+
+    # Download the file from Hugging Face Hub
+    return hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        subfolder=subfolder,
+        local_dir=target_folder_path
+    )
+
+# Example usage:
+if __name__ == "__main__":
+    # Example to download the model from noahcao/sapiens-pose-coco repository
+    target_dir = "xxxx/sapiens/sapiens_lite_host/torchscript/pose/checkpoints/sapiens_2b"
+
+    downloaded_file = download_from_huggingface_hub(
+        target_folder_path=target_dir,
+        repo_id="noahcao/sapiens-pose-coco",
+        filename="sapiens_2b_coco_best_coco_AP_822_torchscript.pt2",
+        subfolder="sapiens_lite_host/torchscript/pose/checkpoints/sapiens_2b"
+    )
+
+    print(f"File downloaded to: {downloaded_file}")
+

deeplabcut/common_tools/json_tools.py

@@ -0,0 +1,29 @@
+import json
+import os
+from pathlib import Path
+
+def convert_json_indent_folder_level(source_folder_path, target_folder_path):
+    """
+    Read JSON files and save them with proper indentation.
+
+    Args:
+        source_folder_path (str): Path to the folder containing original JSON files
+        target_folder_path (str): Path to save the formatted JSON files
+    """
+    # Create target folder if it doesn't exist
+    Path(target_folder_path).mkdir(parents=True, exist_ok=True)
+
+    # Get all JSON files from source folder
+    json_files = [f for f in os.listdir(source_folder_path) if f.endswith('.json')]
+    print(f"Processing {len(json_files)} JSON files...")
+
+    for json_file in json_files:
+        source_path = os.path.join(source_folder_path, json_file)
+        target_path = os.path.join(target_folder_path, json_file)
+
+        # Read source JSON file and save with proper indentation
+        with open(source_path, 'r') as f:
+            data = json.load(f)
+
+        with open(target_path, 'w') as f:
+            json.dump(data, f, indent=4)

deeplabcut/modelzoo/video_inference.py

@@ -12,6 +12,7 @@
 
 import json
 import os
+import torch
 from pathlib import Path
 from typing import Optional, Union
 
@@ -34,6 +35,21 @@
     video_to_frames,
 )
 
+def get_checkpoint_epoch(checkpoint_path):
+    """
+    Load a PyTorch checkpoint and return the current epoch number.
+
+    Args:
+        checkpoint_path (str): Path to the checkpoint file
+
+    Returns: 
+        int: Current epoch number, or None if not found
+    """ 
+    checkpoint = torch.load(checkpoint_path)
+    if 'metadata' in checkpoint and 'epoch' in checkpoint['metadata']:
+        return checkpoint['metadata']['epoch']
+    else:
+        return 0
 
 def get_checkpoint_epoch(checkpoint_path):
     """
@@ -517,7 +533,6 @@ def video_inference_superanimal(
 
             model_snapshot_prefix = f"snapshot-{model_name}"
             detector_snapshot_prefix = f"snapshot-{detector_name}"
-
             config["runner"]["snapshot_prefix"] = model_snapshot_prefix
             config["detector"]["runner"]["snapshot_prefix"] = detector_snapshot_prefix
 
@@ -599,6 +614,7 @@ def video_inference_superanimal(
                     detector_batch_size=video_adapt_batch_size,
                 )
 
+
             # after video adaptation, re-update the adapted checkpoint path, if the checkpoint does not exist, use the best checkpoint
             adapted_detector_checkpoint = (
                 model_folder