Resizing in eval; new inference script

src/deepgraphpose/helpers/csv_to_numpy.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+import argparse
+import pandas as pd
+import numpy as np
+
+"""Convert CSV representation to numpy array
+"""
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="dgp-csv-to-numpy")
+
+    # dataset
+    parser.add_argument("--csv_paths", type=str, nargs="+", required=True, help='Order by views since they will be stacked')
+    parser.add_argument("--output_path", type=str, default="./output.npy")
+    parser.add_argument("--thresh", type=float, default=0.0)
+
+    return parser.parse_args()
+
+
+def count_bodyparts(df):
+    count = 0
+    for col in df.columns:
+        if 'x' in col:
+            count += 1
+
+    return count
+
+
+if __name__ == "__main__":
+    args = get_args()
+    csv_file_list = args.csv_paths
+    thresh = args.thresh
+
+    multi_x_coords = []
+    multi_y_coords = []
+
+    for i, csv_file in enumerate(csv_file_list):
+        csv_file = Path(csv_file).resolve()
+        df = pd.read_csv(csv_file, skiprows=[0, 1])
+        # x, y then x.1, y.1, then x.2, y.2, etc.
+        num_bp = count_bodyparts(df)
+        x_coords = []
+        y_coords = []
+        likelihoods = []
+
+        for j in range(num_bp):
+            if i == 1:
+                j = num_bp - j - 1
+
+            if j == 0:
+                index = ''
+            else:
+                index = '.' + str(j)
+
+            x = df['x' + index].values
+            y = df['y' + index].values
+            lh = df['likelihood' + index].values
+            x_coords.append(x)
+            y_coords.append(y)
+            likelihoods.append(lh)
+
+        x_coords = np.asarray(x_coords)
+        y_coords = np.asarray(y_coords)
+        likelihoods = np.asarray(likelihoods)
+
+        nan_indices = likelihoods <= thresh
+        x_coords[nan_indices] = np.nan
+        y_coords[nan_indices] = np.nan
+
+        multi_x_coords.append(x_coords)
+        multi_y_coords.append(y_coords)
+
+    multi_x_coords = np.asarray(multi_x_coords).transpose(0, 2, 1)[:, :, :, np.newaxis]
+    multi_y_coords = np.asarray(multi_y_coords).transpose(0, 2, 1)[:, :, :, np.newaxis]
+
+    multiview_arr = np.concatenate((multi_x_coords, multi_y_coords), axis=-1)
+
+    np.save(args.output_path, multiview_arr)

src/deepgraphpose/models/eval.py

@@ -11,6 +11,7 @@
 from skimage.util import img_as_ubyte
 from tqdm import tqdm
 from pathlib import Path
+from PIL import Image
 vers = tf.__version__.split('.')
 if int(vers[0]) == 1 and int(vers[1]) > 12:
     TF = tf.compat.v1
@@ -224,7 +225,7 @@ def setup_dgp_eval_graph(dlc_cfg, dgp_model_file, loc_ref=False, gauss_len=1, ga
 
 def estimate_pose(
         proj_cfg_file, dgp_model_file, video_file, output_dir, shuffle=1, save_pose=True,
-        save_str=''):
+        save_str='', new_size=None):
     """Estimate pose on an arbitrary video.
 
     Parameters
@@ -243,6 +244,8 @@ def estimate_pose(
         True to save out pose in csv/hdf5 file
     save_str : str, optional
         additional string to append to labeled video file name
+    new_size : tuple, optional
+        Size for resizing frames (height, width)
 
     Returns
     -------
@@ -279,21 +282,33 @@ def estimate_pose(
     # -------------------
     try:
         dlc_cfg.net_type = 'resnet_50'
-        sess, mu_n, _, _, _, inputs = setup_dgp_eval_graph(dlc_cfg, dgp_model_file)
+        sess, mu_n, softmax_tensor, scmap_tf, _, inputs = setup_dgp_eval_graph(dlc_cfg, dgp_model_file)
     except:
         dlc_cfg.net_type = 'resnet_101'
-        sess, mu_n, _, _, _, inputs = setup_dgp_eval_graph(dlc_cfg, dgp_model_file)
-        
+        sess, mu_n, softmax_tensor, scmap_tf, _, inputs = setup_dgp_eval_graph(dlc_cfg, dgp_model_file)
+
     print('\n')
     pbar = tqdm(total=n_frames, desc='processing video frames')
     markers = np.zeros((n_frames, dlc_cfg.num_joints, 2))
     likelihoods = np.zeros((n_frames, dlc_cfg.num_joints))
+
     for i, frame in enumerate(video_clip.iter_frames()):
+        if new_size is not None:
+            frame = Image.fromarray(frame)
+            scale_x = frame.width / new_size[1]
+            scale_y = frame.height / new_size[0]
+            # width, height
+            frame = frame.resize(size=(new_size[1], new_size[0]))
+            frame = np.asarray(frame)
+        else:
+            scale_x = 1
+            scale_y = 1
+
         # get resnet output
         ff = img_as_ubyte(frame)
-        mu_n_batch = sess.run(mu_n, feed_dict={inputs: ff[None, :, :, :]})
+        mu_n_batch, softmax_tensor_batch, scmap_tf_batch = sess.run([mu_n, softmax_tensor, scmap_tf], feed_dict={inputs: ff[None, :, :, :]})
         markers[i] = mu_n_batch * dlc_cfg.stride + 0.5 * dlc_cfg.stride
-        likelihoods[i] = 0.5
+        likelihoods[i] = np.max(np.max(softmax_tensor_batch, axis=1), axis=1)
 
         pbar.update(1)
 
@@ -305,8 +320,8 @@ def estimate_pose(
     # save labels
     # -------------------
     labels = {
-        'x': markers[:, :, 1],
-        'y': markers[:, :, 0],
+        'x': markers[:, :, 1] * scale_x,
+        'y': markers[:, :, 0] * scale_y,
         'likelihoods': likelihoods}
 
     # convert to DLC-like csv/hdf5

src/deepgraphpose/models/infer_dgp.py

@@ -0,0 +1,85 @@
+# If you have collected labels using DLC's GUI you can run DGP with the following
+"""Main fitting function for DGP.
+   step 0: run DLC
+   step 1: run DGP with labeled frames only
+   step 2: run DGP with spatial clique
+   step 3: do prediction on all videos
+"""
+import argparse
+import os
+from os import listdir
+from os.path import isfile, join
+from pathlib import Path
+import sys
+import yaml
+
+if sys.platform == "darwin":
+    import wx
+
+    if int(wx.__version__[0]) > 3:
+        wx.Thread_IsMain = wx.IsMainThread
+
+os.environ["DLClight"] = "True"
+os.environ["Colab"] = "True"
+from deeplabcut.utils import auxiliaryfunctions
+from deepgraphpose.models.fitdgp import fit_dlc, fit_dgp, fit_dgp_labeledonly
+from deepgraphpose.models.fitdgp_util import get_snapshot_path
+from deepgraphpose.models.eval import plot_dgp, estimate_pose
+from pathlib import Path
+import argparse
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="causal-gen")
+
+    # dataset
+    parser.add_argument("--video_paths", type=str, nargs="+", required=True)
+    parser.add_argument("--output_dir", type=str, default="./outputs")
+    parser.add_argument(
+        "--snapshot",
+        type=str,
+        required=True
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        required=True
+    )
+    parser.add_argument("--dot_sizes", type=int, nargs="+", default=3)
+    parser.add_argument("--new_hw", type=int, nargs="+")
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = get_args()
+    video_file_list = args.video_paths
+    config = Path(args.config).resolve()
+    snapshot = Path(args.snapshot).resolve()
+    output_dir = Path(args.output_dir).resolve()
+    dot_sizes = args.dot_sizes
+    output_dir.mkdir(exist_ok=True, parents=True)
+    if len(dot_sizes) == 1:
+        dot_sizes *= len(video_file_list)
+
+    for i, video_file in enumerate(video_file_list):
+        video_file = Path(video_file)
+
+        estimate_pose(
+            proj_cfg_file=config,
+            dgp_model_file=snapshot,
+            video_file=video_file,
+            output_dir=output_dir,
+            save_str=video_file.stem,
+            new_size=args.new_hw
+        )
+        plot_dgp(
+            video_file=video_file,
+            output_dir=output_dir,
+            label_dir=output_dir,
+            proj_cfg_file=config,
+            dgp_model_file=snapshot,
+            save_str=video_file.stem,
+            dotsize=dot_sizes[i]
+        )
+