from transID to master (#1) (DeepLabCut#1764)

Added support for transformer-based ReID

Author: yeshaokai

.gitignore

@@ -9,7 +9,7 @@
 .local
 .DS_Store
 examples/.DS_Store
-
+*~
 # Tensorflow checkpoints
 *.ckpt
 snapshot-*

deeplabcut/__init__.py

@@ -85,6 +85,7 @@
     evaluate_network,
     return_evaluate_network_data,
     analyze_videos,
+    create_tracking_dataset,
     analyze_time_lapse_frames,
     convert_detections2tracklets,
     extract_maps,
@@ -95,6 +96,7 @@
     export_model,
 )
 
+
 from deeplabcut.pose_estimation_3d import (
     calibrate_cameras,
     check_undistortion,
@@ -111,4 +113,6 @@
 from deeplabcut.post_processing import filterpredictions, analyzeskeleton
 
 
+
+
 from deeplabcut.version import __version__, VERSION

deeplabcut/create_project/new.py

@@ -92,7 +92,7 @@ def create_new_project(
     # Create project and sub-directories
     if not DEBUG and project_path.exists():
         print('Project "{}" already exists!'.format(project_path))
-        return
+        return  os.path.join(str(project_path), "config.yaml")
     video_path = project_path / "videos"
     data_path = project_path / "labeled-data"
     shuffles_path = project_path / "training-datasets"

deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py

@@ -34,7 +34,10 @@
 )
 
 def format_multianimal_training_data(
-    df, train_inds, project_path, n_decimals=2,
+    df,
+    train_inds,
+    project_path,
+    n_decimals=2,
 ):
     train_data = []
     nrows = df.shape[0]
@@ -318,7 +321,10 @@ def create_multianimaltraining_dataset(
 
         # Make training file!
         data = format_multianimal_training_data(
-            Data, trainIndices, cfg["project_path"], numdigits,
+            Data,
+            trainIndices,
+            cfg["project_path"],
+            numdigits,
         )
 
         if len(trainIndices) > 0:
@@ -374,7 +380,10 @@ def create_multianimaltraining_dataset(
             )
             path_test_config = str(
                 os.path.join(
-                    cfg["project_path"], Path(modelfoldername), "test", "pose_cfg.yaml",
+                    cfg["project_path"],
+                    Path(modelfoldername),
+                    "test",
+                    "pose_cfg.yaml",
                 )
             )
             path_inference_config = str(
@@ -475,7 +484,10 @@ def create_multianimaltraining_dataset(
 
 
 def convert_cropped_to_standard_dataset(
-    config_path, recreate_datasets=True, delete_crops=True, back_up=True,
+    config_path,
+    recreate_datasets=True,
+    delete_crops=True,
+    back_up=True,
 ):
     import pandas as pd
     import pickle
@@ -514,7 +526,8 @@ def convert_cropped_to_standard_dataset(
         return
 
     datasets_folder = os.path.join(
-        project_path, auxiliaryfunctions.GetTrainingSetFolder(cfg),
+        project_path,
+        auxiliaryfunctions.GetTrainingSetFolder(cfg),
     )
     df_old = pd.read_hdf(
         os.path.join(datasets_folder, "CollectedData_" + cfg["scorer"] + ".h5"),

deeplabcut/generate_training_dataset/trainingsetmanipulation.py

@@ -245,14 +245,14 @@ def dropimagesduetolackofannotation(config):
 
 def dropunlabeledframes(config):
     """
-    Drop entries such that all the bodyparts are not labeled from the annotation files, i.e. h5 and csv files 
+    Drop entries such that all the bodyparts are not labeled from the annotation files, i.e. h5 and csv files
     Will be carried out iteratively for all *folders* in labeled-data.
-    
+
     Parameter
     ----------
     config : string
         String containing the full path of the config file in the project.
-    
+
     """
     cfg = auxiliaryfunctions.read_config(config)
     videos = cfg["video_sets"].keys()
@@ -275,9 +275,9 @@ def dropunlabeledframes(config):
             DC.to_csv(
                 os.path.join(str(folder), "CollectedData_" + cfg["scorer"] + ".csv")
             )
-            
+
             print("Dropped ", dropped, "entries in ",folder)
-    
+
     print("Done.")
 
 def check_labels(
@@ -747,7 +747,7 @@ def create_training_dataset(
 
     augmenter_type: string
         Type of augmenter. Currently default, imgaug, tensorpack, and deterministic are supported.
-        
+
     posecfg_template: string (optional, default=None)
         Path to a pose_cfg.yaml file to use as a template for generating the new one for the current iteration. Useful if you
         would like to start with the same parameters a previous training iteration. None uses the default pose_cfg.yaml.
@@ -834,7 +834,7 @@ def create_training_dataset(
             "deterministic",
         ]:
             raise ValueError("Invalid augmenter type:", augmenter_type)
-        
+
         if posecfg_template:
             if net_type != prior_cfg["net_type"]:
                 print(
@@ -844,7 +844,7 @@ def create_training_dataset(
                 print(
                     "WARNING: Specified augmenter_type does not match dataset_type from posecfg_template path entered. Proceed with caution."
                 )
-                
+
         # Loading the encoder (if necessary downloading from TF)
         dlcparent_path = auxiliaryfunctions.get_deeplabcut_path()
         if not posecfg_template:

deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py

@@ -296,6 +296,7 @@ def evaluate_multianimal_full(
                     )
 
                     data_path = resultsfilename.split(".h5")[0] + "_full.pickle"
+
                     if plotting:
                         foldername = os.path.join(
                             str(evaluationfolder),
@@ -308,7 +309,12 @@ def evaluate_multianimal_full(
                     if os.path.isfile(data_path):
                         print("Model already evaluated.", resultsfilename)
                     else:
-                        sess, inputs, outputs = predict.setup_pose_prediction(dlc_cfg)
+
+                        (
+                            sess,
+                            inputs,
+                            outputs,
+                        ) = predict.setup_pose_prediction(dlc_cfg)
 
                         PredicteData = {}
                         dist = np.full((len(Data), len(all_bpts)), np.nan)
@@ -359,6 +365,7 @@ def evaluate_multianimal_full(
                                 :, ["sample", "y", "x", "bodyparts"]
                             ].to_numpy()
                             peaks_gt[:, 1:3] = (peaks_gt[:, 1:3] - stride // 2) / stride
+
                             pred = predictma.predict_batched_peaks_and_costs(
                                 dlc_cfg,
                                 np.expand_dims(frame, axis=0),
@@ -367,6 +374,7 @@ def evaluate_multianimal_full(
                                 outputs,
                                 peaks_gt.astype(int),
                             )
+                            
                             if not pred:
                                 continue
                             else:

deeplabcut/pose_estimation_tensorflow/core/predict.py

@@ -28,12 +28,13 @@
 from .openvino.session import OpenVINOSession
 
 
-def setup_pose_prediction(cfg, allow_growth=False):
+def setup_pose_prediction(cfg, allow_growth=False, collect_extra=False):
     tf.compat.v1.reset_default_graph()
     inputs = tf.compat.v1.placeholder(
         tf.float32, shape=[cfg["batch_size"], None, None, 3]
     )
     net_heads = PoseNetFactory.create(cfg).test(inputs)
+    extra_dict = {}
     outputs = [net_heads["part_prob"]]
     if cfg["location_refinement"]:
         outputs.append(net_heads["locref"])
@@ -44,6 +45,9 @@ def setup_pose_prediction(cfg, allow_growth=False):
 
     outputs.append(net_heads["peak_inds"])
 
+    if collect_extra:
+        extra_dict["features"] = net_heads["features"]
+
     restorer = tf.compat.v1.train.Saver()
 
     if allow_growth:
@@ -58,7 +62,10 @@ def setup_pose_prediction(cfg, allow_growth=False):
     # Restore variables from disk.
     restorer.restore(sess, cfg["init_weights"])
 
-    return sess, inputs, outputs
+    if collect_extra:
+        return sess, inputs, outputs, extra_dict
+    else:
+        return sess, inputs, outputs
 
 
 def extract_cnn_output(outputs_np, cfg):

deeplabcut/pose_estimation_tensorflow/core/predict_multianimal.py

@@ -211,7 +211,14 @@ def predict_batched_peaks_and_costs(
     peaks_gt=None,
     n_points=10,
     n_decimals=3,
+    extra_dict=None,
 ):
+
+    if extra_dict:
+        features = sess.run(extra_dict['features'],
+                            feed_dict = {inputs:images_batch}
+                            )
+    
     scmaps, locrefs, *pafs, peaks = sess.run(outputs, feed_dict={inputs: images_batch})
     if ~np.any(peaks):
         return []
@@ -249,8 +256,10 @@ def predict_batched_peaks_and_costs(
         )
         for i, costs in enumerate(costs_gt):
             preds[i]["groundtruth_costs"] = costs
-
-    return preds
+    if extra_dict:
+        return preds, features
+    else:
+        return preds
 
 
 def find_local_maxima(scmap, radius, threshold):

deeplabcut/pose_estimation_tensorflow/core/train.py

@@ -148,7 +148,7 @@ def train(
     maxiters,
     max_to_keep=5,
     keepdeconvweights=True,
-    allow_growth=False,
+    allow_growth=True,
 ):
     start_path = os.getcwd()
     os.chdir(

deeplabcut/pose_estimation_tensorflow/core/train_multianimal.py

@@ -38,7 +38,7 @@ def train(
     maxiters,
     max_to_keep=5,
     keepdeconvweights=True,
-    allow_growth=False,
+    allow_growth=True,
 ):
     start_path = os.getcwd()
     os.chdir(