Safer DLC version retrieval (#18)

* better upstream version grab; adding some descriptions * Update dlc2nwb/utils.py Co-authored-by: Jessy Lauer <[email protected]> * Update utils.py Co-authored-by: Jessy Lauer <[email protected]>
DeepLabCut · Jan 26, 2023 · 9e200b8 · 9e200b8
1 parent a2c0a29
commit 9e200b8
Showing 1 changed file with 27 additions and 9 deletions.
diff --git a/dlc2nwb/utils.py b/dlc2nwb/utils.py
@@ -1,24 +1,34 @@
-import cv2
 import datetime
 import os
+import pickle
+import warnings
 from pathlib import Path
+from platform import python_version
+
+import cv2
+import yaml
 import numpy as np
 import pandas as pd
-import pickle
-import warnings
 from hdmf.build.warnings import DtypeConversionWarning
+from packaging.version import Version # Installed with setuptools
 from pynwb import NWBFile, NWBHDF5IO
 from ndx_pose import PoseEstimationSeries, PoseEstimation
-import yaml
 from ruamel.yaml import YAML
 
 # If available determine version
 try:
- from deeplabcut import __version__
- deeplabcut_version = __version__ 
+ if Version(python_version()) >= Version("3.8"):
+ from importlib.metadata import version as get_package_version # Only available in Python>=3.8
+
+ deeplabcut_version = get_package_version("deeplabcut")
+ else:
+ from pkg_resources import get_distribution # Installed with setuptools
+
+ deeplabcut_version = get_distribution("deeplabcut").version
 except ModuleNotFoundError:
  deeplabcut_version = None
 
+
 def read_config(configname):
  """
  Reads structured config file defining a project.
@@ -176,8 +186,16 @@ def _get_pes_args(config_file, h5file, individual_name, infer_timestamps=True):
  return scorer, df, video, paf_graph, timestamps, cfg
 
 
-def _write_pes_to_nwbfile(nwbfile, animal, df_animal, scorer, video, paf_graph, timestamps,
- exclude_nans):
+def _write_pes_to_nwbfile(
+ nwbfile,
+ animal,
+ df_animal,
+ scorer,
+ video, # Expects this to be a tuple; first index is string path, second is the image shape as "0, width, 0, height"
+ paf_graph,
+ timestamps,
+ exclude_nans,
+): 
  pose_estimation_series = []
  for kpt, xyp in df_animal.groupby(level="bodyparts", axis=1, sort=False):
  data = xyp.to_numpy()
@@ -200,7 +218,7 @@ def _write_pes_to_nwbfile(nwbfile, animal, df_animal, scorer, video, paf_graph,
  confidence_definition="Softmax output of the deep neural network.",
  )
  pose_estimation_series.append(pes)
- 
+
  pe = PoseEstimation(
  pose_estimation_series=pose_estimation_series,
  description="2D keypoint coordinates estimated using DeepLabCut.",