diff --git a/src/deepgraphpose/helpers/csv_to_numpy.py b/src/deepgraphpose/helpers/csv_to_numpy.py new file mode 100644 index 0000000..24e36ab --- /dev/null +++ b/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) diff --git a/src/deepgraphpose/models/eval.py b/src/deepgraphpose/models/eval.py index ba4cb83..972303d 100644 --- a/src/deepgraphpose/models/eval.py +++ b/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 diff --git a/src/deepgraphpose/models/infer_dgp.py b/src/deepgraphpose/models/infer_dgp.py new file mode 100644 index 0000000..09c1735 --- /dev/null +++ b/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] + ) +