def process_video(args):
file, out_action_video, video = args
out_action_video_file = osp.join(
out_action_video, osp.basename(file).split(".")[0]
)
file_utils.mkdir(out_action_video_file)
tstamps = []
for i, (frame, tstamp) in enumerate(video_to_frames(file, fps=0)):
file_utils.write_image_to_jpeg(
frame, osp.join(out_action_video_file, "{}.jpg".format(i))
)
tstamps.append(tstamp)
file_utils.write_timestamps_to_txt(
tstamps, osp.join(out_action_video_file, "img_timestamps.txt")
)
def main(args):
plt.switch_backend("Agg")
emb_files = file_utils.get_files(osp.join(args.embs_path, "embs"),
"*.npy",
sort=False)
print("Found files: {}".format(emb_files))
for i, emb_file in enumerate(emb_files):
embs, frames = [], []
with open(emb_file, "rb") as f:
query_dict = np.load(f, allow_pickle=True).item()
for j in range(len(query_dict["embs"])):
curr_embs = query_dict["embs"][j]
if args.l2_normalize:
curr_embs = [x / (np.linalg.norm(x) + 1e-7) for x in curr_embs]
embs.append(curr_embs)
frames.append(query_dict["frames"][j])
# generate video name
video_path = osp.join(args.embs_path, "videos")
file_utils.mkdir(video_path)
ext = ".mp4" if args.align else "_original.mp4"
video_path = osp.join(video_path,
osp.basename(emb_file).split(".")[0] + ext)
if not args.align:
print("Playing videos without alignment.")
create_original_videos(frames, video_path, args.interval)
else:
print("Aligning videos.")
create_video(
embs,
frames,
video_path,
args.use_dtw,
query=args.reference_video,
candidate=args.candidate_video,
interval=args.interval,
)
def videos_to_images(in_dir, out_dir, num_cores):
set_start_method("spawn")
with Pool(num_cores) as pool:
action_dirs = file_utils.get_subdirs(in_dir)
for action_dir in action_dirs:
print("Processing ", osp.basename(action_dir))
out_action = osp.join(out_dir, osp.basename(action_dir))
file_utils.mkdir(out_action)
videos = file_utils.get_subdirs(action_dir)
for video in videos:
out_action_video = osp.join(out_action, osp.basename(video))
file_utils.mkdir(out_action_video)
files = file_utils.get_files(video, "*.mp4", False)
func_args = [[f, out_action_video, video] for f in files]
for _ in tqdm(
pool.imap_unordered(process_video, func_args),
total=len(files),
):
pass
file_utils.copy_file(
osp.join(video, "joint_angles.txt"),
osp.join(out_action_video, "joint_angles.txt"),
)
def __init__(self, checkpoint, directory, device, max_to_keep=10):
"""Constructor.
Args:
checkpoint (Checkpoint): An instance of `Checkpoint`.
directory (str): The directory in which checkpoints will be saved.
device (torch.device): The computing device on which to restore
checkpoints.
max_to_keep (int): The maximum number of checkpoints to keep.
Amongst all saved checkpoints, checkpoints will be deleted
oldest first, until `max_to_keep` remain.
"""
assert max_to_keep > 0, "max_to_keep should be a positive integer."
self.checkpoint = checkpoint
self.directory = directory
self.max_to_keep = max_to_keep
self.device = device
self.latest_checkpoint = None
# create checkpoint directory if it doesn't
# already exist
mkdir(self.directory)
def init_experiment(
logdir, config, config_file=None, override_list=None, transient=False,
):
"""Initializes a training experiment.
Instantiates the compute device (CPU, GPU), serializes the config
file to the log directory, optionally updates the config variables
with values from a provided yaml file and seeds the RNGs.
Args:
logdir (str): Path to the log directory.
config (dict): The module-wide config dict.
config_file (str): Path to an experimental run yaml
file containing config options we'd like to override.
override_list (list): Additional command line args we'd
like to override in the config file.
transient (bool): Set to `True` to make a transient session,
i.e. a session where the logging and config params are
not saved to disk. This is useful for debugging sessions.
"""
# init compute device
if torch.cuda.is_available():
device = torch.device("cuda")
logging.info(
"Using GPU {}.".format(torch.cuda.get_device_name(device))
)
else:
logging.info("No GPU found. Falling back to CPU.")
device = torch.device("cpu")
# create logdir and update config dict
if not transient:
# if a yaml file already exists in the
# log directory, it means we're resuming
# from a previous run.
# so we update the values of our config
# file with the values in the yaml file.
if osp.exists(osp.join(logdir, "config.yml")):
logging.info(
"Config yaml already exists in log dir. Resuming training."
)
update_config(
config, osp.join(logdir, "config.yml"), override_list
)
# if no yaml file exists in the log directory,
# it means we're starting a new experiment.
# so we want to update our config file
# but also serialize it to the log dir.
else:
# create the log directory if
# it doesn't already exist
mkdir(logdir)
update_config(
config,
config_file,
override_list,
osp.join(logdir, "config.yml"),
)
else:
logging.info("Transient model turned ON.")
update_config(config, config_file, override_list)
# seed rngs
if config.SEED is not None:
logging.info(f"Experiment seed: {config.SEED}.")
seed_rng(config.SEED)
else:
logging.info("No RNG seed has been set for this experiment.")
return config, device
def main(args):
if torch.cuda.is_available():
device = torch.device("cuda")
logging.info("Using GPU {}.".format(
torch.cuda.get_device_name(device)))
else:
logging.info("No GPU found. Falling back to CPU.")
device = torch.device("cpu")
# initialize experiment
opts = [
"SAMPLING.STRIDE_ALL_SAMPLER",
args.stride,
]
config_path = osp.join(args.logdir, "config.yml")
config, device = experiment_utils.init_experiment(
args.logdir,
CONFIG,
config_path,
opts,
)
# load model and data loaders
debug = {
"sample_sequential": not args.shuffle,
"augment": False,
"labeled": None, # "both" if args.negative else "pos",
}
model, _, loaders, _, _ = experiment_utils.get_factories(config,
device,
debug=debug)
loaders = (loaders["downstream_valid"]
if args.split == "valid" else loaders["downstream_train"])
# load model checkpoint
if args.model_ckpt is not None:
checkpoint.Checkpoint(model).restore(args.model_ckpt, device)
else:
checkpoint.CheckpointManager.load_latest_checkpoint(
checkpoint.Checkpoint(model),
osp.join(config.DIRS.CKPT_DIR,
osp.basename(osp.normpath(args.logdir))),
device,
)
model.to(device).eval()
# figure out max batch size that's
# a multiple of the number of context
# frames.
# this is so we can support large videos
# with many frames.
lcm = model.num_ctx_frames
max_batch_size = math.floor(128 / lcm) * lcm
# create save folder
save_path = osp.join(
config.DIRS.DIR,
args.save_path,
osp.basename(osp.normpath(args.logdir)),
"embs",
)
file_utils.mkdir(save_path)
# iterate over every class action
pbar = tqdm.tqdm(loaders.items(), leave=False)
for action_name, loader in pbar:
msg = "embedding {}".format(action_name)
pbar.set_description(msg)
(
embeddings,
seq_lens,
steps,
vid_frames,
names,
labels,
phase_labels,
) = ([] for i in range(7))
for batch_idx, batch in enumerate(loader):
if args.max_embs != -1 and batch_idx >= args.max_embs:
break
# unpack batch data
frames = batch["frames"]
chosen_steps = batch["frame_idxs"].to(device)
seq_len = batch["video_len"].to(device)
name = batch["video_name"][0]
# label = batch["success"][0]
# phase_label = None
# if "phase_labels" in batch:
# phase_label = batch["phase_labels"].to(device)
# forward through model to compute embeddings
with torch.no_grad():
if frames.shape[1] > max_batch_size:
embs = []
for i in range(math.ceil(frames.shape[1] /
max_batch_size)):
sub_frames = frames[:, i * max_batch_size:(i + 1) *
max_batch_size].to(device)
sub_embs = model(sub_frames)["embs"]
embs.append(sub_embs.cpu())
embs = torch.cat(embs, dim=1)
else:
embs = model(frames.to(device))["embs"]
# store
embeddings.append(embs.cpu().squeeze().numpy())
seq_lens.append(seq_len.cpu().squeeze().numpy())
steps.append(chosen_steps.cpu().squeeze().numpy())
# if phase_label is not None:
# phase_labels.append(phase_label.cpu().squeeze().numpy())
names.append(name)
# labels.append(label.item())
if args.keep_frames:
frames = frames[0]
frames = UnNormalize()(frames)
frames = frames.view(
embs.shape[1],
config.SAMPLING.NUM_CONTEXT_FRAMES,
*frames.shape[1:],
)
vid_frames.append(frames.cpu().squeeze().numpy()[:,
-1].transpose(
0, 2, 3,
1))
data = {
"embs": embeddings,
"seq_lens": seq_lens,
"steps": steps,
"names": names,
# "labels": labels,
}
if args.keep_frames:
data["frames"] = vid_frames
# if phase_labels:
# data["phase_labels"] = phase_labels
np.save(osp.join(save_path, action_name), data)
INCLUDE_DEPTH = False
in_dir = "/home/kevin/repos/kronos/kronos/data/mime/"
out_dir = "/home/kevin/repos/kronos/kronos/data/mime_processed/"
for folder_name, folder_id in mime_map.items():
in_pouring_dir = osp.join(in_dir, "{}".format(folder_id))
wanted_folders = ["hd_kinect_rgb", "rd_kinect_rgb"]
if not OVERHEAD_ONLY:
wanted_folders += ["rd_kinect_rgb"]
demonstrations = file_utils.get_subdirs(in_pouring_dir)
num_valid = int(VALIDATION_FRAC * len(demonstrations))
valid_demonstrations = demonstrations[:num_valid]
train_demonstrations = demonstrations[num_valid:]
for name, demonstrations in zip(
["train", "valid"], [train_demonstrations, valid_demonstrations]
):
out_pouring_dir = osp.join(
out_dir, name, folder_name.lower().replace(" ", "_")
)
file_utils.mkdir(out_pouring_dir)
cnter = 0
for in_demo in tqdm(demonstrations):
left_or_right = (
[] if OVERHEAD_ONLY else load_joint_angles(in_demo)
)
for wanted in wanted_folders + left_or_right:
in_wanted = osp.join(in_demo, wanted)
out_wanted = osp.join(out_pouring_dir, "{}".format(cnter))
if not osp.isdir(out_wanted):
copy_tree(in_wanted, out_wanted)
cnter += 1
# "--resume",
# "True",
# ]
# )
# train lstm model
out = subprocess.check_output([
"python",
"scripts/train_lstm.py",
"--logdir",
"kronos/logs/{}".format(experiment_name),
"--query",
query_action,
"--l2_normalize",
"False",
"--batch_size",
"20",
"--stride",
"20",
"--max_iters",
"300",
])
best_acc = float(out.decode().split("\n")[-2].split(":")[-1].lstrip())
best_accs[query_action] = best_acc
print("{}: {}".format(query_action, best_acc))
save_dir = f"tmp/exp_1/{task}"
file_utils.mkdir(save_dir)
file_utils.write_json(osp.join(save_dir, "result.json"), best_accs)
" By default, stored in the input directory.",
)
parser.add_argument(
"--train_frac",
type=float,
default=0.975,
help="The fraction used for train.",
)
args, unparsed = parser.parse_known_args()
# Get list of subdirectories in the input dir
input_subdirs = file_utils.get_subdirs(args.input_dir)
# Create output dir if it does not exist.
if args.output_dir is not None:
file_utils.mkdir(args.output_dir)
output_dir = args.output_dir
operator = file_utils.copy_folder
else:
output_dir = args.input_dir
operator = file_utils.move_folder
base_train_dir = osp.join(output_dir, "train")
base_valid_dir = osp.join(output_dir, "valid")
for input_subdir in input_subdirs:
video_dirs = file_utils.get_subdirs(input_subdir)
# Create train and validation splits.
np.random.shuffle(video_dirs)
num_train = int(args.train_frac * len(video_dirs))
"--width",
type=int,
default=256,
help="The width of processed video frames if resize is set to True.",
)
parser.add_argument(
"--num_cores",
type=int,
default=cpu_count(),
help="The number of cores to use in parallel.",
)
args, unparsed = parser.parse_known_args()
# create output dir if it doesn't exist
if args.output_dir is not None:
file_utils.mkdir(args.output_dir)
else:
args.output_dir = args.input_dir
frame_dir = osp.join(args.input_dir, "frames")
label_dir = osp.join(args.input_dir, "labels")
if args.resize:
args.resize = (args.height, args.width)
logging.info("Resizing images to ({}, {}).".format(*args.resize))
else:
args.resize = None
# read all label files and figure out which
# videos belong to which class
label_files = file_utils.get_files(label_dir, "*.mat")
label_classes = [loadmat(lab)["action"][0] for lab in label_files]