138 lines
4.0 KiB
Python
138 lines
4.0 KiB
Python
import os
|
|
import glob
|
|
import pandas as pd
|
|
from copy import deepcopy
|
|
import warnings
|
|
|
|
|
|
def update_args(args):
|
|
args = deepcopy(args)
|
|
if args.type == "egobody":
|
|
if args.root is None:
|
|
warnings.warn(
|
|
"args.root flag is set to None, using root directory in source"
|
|
)
|
|
args.root = "/path/to/egobody"
|
|
if args.seqs is None:
|
|
args.seqs = get_egobody_seqs(args.root, args.split)
|
|
args.img_dirs = [
|
|
glob.glob(f"{args.root}/egocentric_color/{seq}/**/PV")[0]
|
|
for seq in args.seqs
|
|
]
|
|
return args
|
|
|
|
if args.type == "3dpw":
|
|
if args.root is None:
|
|
warnings.warn(
|
|
"args.root flag is set to None, using root directory in source"
|
|
)
|
|
args.root = "/path/to/3DPW"
|
|
if args.split == "val":
|
|
args.split = "validation"
|
|
if args.seqs is None:
|
|
args.seqs = get_3dpw_seqs(args.root, args.split)
|
|
args.img_name = "imageFiles"
|
|
|
|
elif args.type == "posetrack":
|
|
if args.root is None:
|
|
warnings.warn(
|
|
"args.root flag is set to None, using root directory in source"
|
|
)
|
|
args.root = "/path/to/posetrack/posetrack2018/posetrack_data"
|
|
args.img_name = f"images/{args.split}"
|
|
|
|
elif args.type == "davis":
|
|
if args.root is None:
|
|
warnings.warn(
|
|
"args.root flag is set to None, using root directory in source"
|
|
)
|
|
args.root = "/path/to/DAVIS"
|
|
args.img_name = "JPEGImages/Full-Resolution"
|
|
args.split = ""
|
|
|
|
else:
|
|
assert args.root is not None
|
|
args.img_name = args.img_name if args.img_name is not None else "images"
|
|
args.split = ""
|
|
|
|
args.img_dirs, args.seqs = get_all_img_dirs(
|
|
f"{args.root}/{args.img_name}", args.seqs
|
|
)
|
|
return args
|
|
|
|
|
|
def isimage(f):
|
|
ext = os.path.splitext(f)[-1].lower()
|
|
return ext in [".png", ".jpg"]
|
|
|
|
|
|
def get_all_img_dirs(data_root, seqs=None):
|
|
"""
|
|
returns all image directories in root, recursively searching, and the sequence names
|
|
"""
|
|
seqs = seqs if seqs is not None else []
|
|
img_dirs = []
|
|
|
|
def get_name(p):
|
|
return p.removeprefix(f"{data_root}/")
|
|
|
|
for root, _, files in os.walk(data_root):
|
|
num_imgs = len([f for f in files if isimage(f)])
|
|
if num_imgs > 0:
|
|
name = get_name(root)
|
|
if len(seqs) == 0 or name in seqs:
|
|
img_dirs.append(root)
|
|
img_dirs = sorted(img_dirs)
|
|
seqs = [get_name(d) for d in img_dirs]
|
|
return img_dirs, seqs
|
|
|
|
|
|
def get_custom_seqs(data_root):
|
|
img_dir = f"{data_root}/images"
|
|
if not os.path.isdir(img_dir):
|
|
return []
|
|
return sorted(os.listdir(img_dir))
|
|
|
|
|
|
def get_egobody_seqs(data_root, split):
|
|
split_file = f"{data_root}/data_splits.csv"
|
|
df = pd.read_csv(split_file)
|
|
if split not in df.columns:
|
|
print(f"{split} not in {split_file}")
|
|
return []
|
|
return sorted(df[split].dropna().tolist())
|
|
|
|
|
|
def get_3dpw_seqs(data_root, split):
|
|
split_dir = f"{data_root}/sequenceFiles/{split}"
|
|
if not os.path.isdir(split_dir):
|
|
return []
|
|
seq_files = sorted(os.listdir(split_dir))
|
|
return [os.path.splitext(f)[0] for f in seq_files]
|
|
|
|
|
|
def get_posetrack_seqs(data_root, split):
|
|
split_dir = f"{data_root}/images/{split}"
|
|
if not os.path.isdir(split_dir):
|
|
return []
|
|
return sorted(os.listdir(split_dir))
|
|
|
|
|
|
def get_davis_seqs(data_root):
|
|
img_root = f"{data_root}/JPEGImages/Full-Resolution"
|
|
if not os.path.isdir(img_root):
|
|
return []
|
|
return sorted(os.listdir(img_root))
|
|
|
|
|
|
def get_img_dir(data_type, data_root, seq, img_name):
|
|
if data_type == "posetrack":
|
|
return f"{data_root}/{img_name}/{seq}"
|
|
if data_type == "3dpw":
|
|
return f"{data_root}/imageFiles/{seq}"
|
|
if data_type == "davis":
|
|
return f"{data_root}/JPEGImages/Full-Resolution/{seq}"
|
|
if data_type == "custom":
|
|
return f"{data_root}/images/{seq}" # custom sequence
|
|
return None
|