def main(out_dir=None,
data_dir=None,
feature_fn_format='feature-seq.pkl',
label_fn_format='label_seq.pkl',
cv_params={}):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
# Load data
trial_ids = utils.getUniqueIds(data_dir,
prefix='trial=',
suffix=feature_fn_format,
to_array=True)
dataset = utils.CvDataset(trial_ids,
data_dir,
feature_fn_format=feature_fn_format,
label_fn_format=label_fn_format,
vocab=[])
utils.saveMetadata(dataset.metadata, out_data_dir)
# Make folds
dataset_size = len(trial_ids)
cv_folds = utils.makeDataSplits(dataset_size,
metadata=dataset.metadata,
**cv_params)
save_cv_folds(cv_folds, os.path.join(out_data_dir, 'cv-folds.json'))
# Check folds
for cv_index, cv_fold in enumerate(cv_folds):
train_data, val_data, test_data = dataset.getFold(cv_fold)
train_feats, train_labels, train_ids = train_data
test_feats, test_labels, test_ids = test_data
val_feats, val_labels, val_ids = val_data
logger.info(
f'CV fold {cv_index + 1} / {len(cv_folds)}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(val_ids)} val, {len(test_ids)} test)'
)
def main(out_dir=None, data_dir=None, labels_dir=None):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
labels_dir = os.path.expanduser(labels_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
filenames = [
utils.stripExtension(fn)
for fn in glob.glob(os.path.join(labels_dir, '*.csv'))
]
metadata = utils.loadMetadata(data_dir)
metadata['seq_id'] = metadata.index
metadata = metadata.set_index(
'dir_name', drop=False).loc[filenames].set_index('seq_id')
seq_ids = np.sort(metadata.index.to_numpy())
logger.info(f"Loaded {len(seq_ids)} sequences from {labels_dir}")
vocab = []
for i, seq_id in enumerate(seq_ids):
seq_id_str = f"seq={seq_id}"
seq_dir_name = metadata['dir_name'].loc[seq_id]
labels_fn = os.path.join(labels_dir, f'{seq_dir_name}.csv')
event_labels = utils.loadVariable(f'{seq_id_str}_labels', data_dir)
assembly_actions = pd.read_csv(labels_fn)
label_seq = parseActions(assembly_actions, event_labels.shape[0],
vocab)
utils.saveVariable(label_seq, f'{seq_id_str}_label-seq', out_data_dir)
plotLabels(os.path.join(fig_dir, f'{seq_id_str}_labels.png'),
label_seq)
writeLabels(os.path.join(fig_dir, f'{seq_id_str}_labels.csv'),
label_seq, vocab)
utils.saveMetadata(metadata, out_data_dir)
utils.saveVariable(vocab, 'vocab', out_data_dir)
def main(out_dir=None, results_file=None, sweep_param_name=None):
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
results_file = os.path.expanduser(results_file)
# Set NUMEXPR_MAX_THREADS to avoid warnings
os.environ["NUMEXPR_MAX_THREADS"] = '10'
data = pd.read_csv(results_file)
if sweep_param_name is None:
means = data.mean() * 100
stds = data.std() * 100
for name in data.columns:
if name == 'loss':
continue
logger.info(f"{name}: {means[name]:.1f}% +/- {stds[name]:.1f}%")
make_hist(os.path.join(out_dir, f"{name}.png"),
data[name].to_numpy())
else:
sweep_vals = data[sweep_param_name].unique()
def gen_rows(sweep_vals):
for val in sweep_vals:
matches_val = data[sweep_param_name] == sweep_vals
matching_data = data.iloc[matches_val]
means = matching_data.mean().rename(
columns=lambda x: f'{x}_mean')
stds = matching_data.std().rename(columns=lambda x: f'{x}_std')
yield pd.concat([means, stds], axis=1)
results = pd.concat([gen_rows(sweep_vals)], axis=0)
results.to_csv(os.path.join(out_dir, "results_aggregate.csv"))
def main(out_dir=None,
video_data_dir=None,
features_dir=None,
activity_labels_dir=None,
gt_keyframes_dir=None,
use_gt_activity_labels=False):
out_dir = os.path.expanduser(out_dir)
video_data_dir = os.path.expanduser(video_data_dir)
features_dir = os.path.expanduser(features_dir)
activity_labels_dir = os.path.expanduser(activity_labels_dir)
if gt_keyframes_dir is not None:
gt_keyframes_dir = os.path.expanduser(gt_keyframes_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
def loadFromDir(var_name, dir_name):
return joblib.load(os.path.join(dir_name, f"{var_name}.pkl"))
def saveToWorkingDir(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f"{var_name}.pkl"))
trial_ids = utils.getUniqueIds(features_dir,
prefix='trial=',
suffix='.pkl')
seq_lens = []
for seq_idx, trial_id in enumerate(trial_ids):
logger.info(
f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})"
)
logger.info(f" Loading data...")
feature_seq = loadFromDir(f"trial={trial_id}_feature-seq",
features_dir)
raw_labels = loadFromDir(f"trial-{trial_id}_action-seq",
video_data_dir)
timestamp_seq = loadFromDir(
f"trial-{trial_id}_rgb-frame-timestamp-seq", video_data_dir)
if use_gt_activity_labels:
activity_label_fn = f"trial={trial_id}_true-label-seq"
else:
activity_label_fn = f"trial={trial_id}_pred-label-seq"
activity_labels = loadFromDir(activity_label_fn, activity_labels_dir)
action_labels = makeActionLabels(raw_labels,
seq_len=feature_seq.shape[0])
if timestamp_seq.shape[0] != feature_seq.shape[0]:
logger.warning(
f"Video dimensions don't match: "
f"{feature_seq.shape} scores, {timestamp_seq.shape} timestamps"
)
continue
# trim sequences
is_activity = activity_labels == 1
if not is_activity.any():
logger.warning(f"No activity detected: skipping trial")
continue
action_labels = action_labels[is_activity, ...]
timestamp_seq = timestamp_seq[is_activity, ...]
feature_seq = feature_seq[is_activity, ...]
if feature_seq.shape[0] != action_labels.shape[0]:
err_str = (
"Data dimensions don't match: "
f"{feature_seq.shape} features, {action_labels.shape} labels")
raise AssertionError(err_str)
seq_lens.append(feature_seq.shape[0])
logger.info(f" Saving output...")
trial_str = f"trial={trial_id}"
saveToWorkingDir(timestamp_seq, f'{trial_str}_timestamp-seq')
saveToWorkingDir(feature_seq, f'{trial_str}_feature-seq')
saveToWorkingDir(action_labels, f'{trial_str}_label-seq')
fn = os.path.join(fig_dir, f"trial={trial_id}.png")
utils.plot_array(feature_seq.T, (action_labels, ), ('action', ), fn=fn)
logger.info(f"min seq len: {min(seq_lens)} max seq len: {max(seq_lens)}")
def main(out_dir=None,
data_dir=None,
annotation_dir=None,
frames_dir=None,
col_format='standard',
win_params={},
slowfast_csv_params={},
label_types=('event', 'action', 'part')):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
annotation_dir = os.path.expanduser(annotation_dir)
frames_dir = os.path.expanduser(frames_dir)
annotation_dir = os.path.join(annotation_dir, 'action_annotations')
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_labels_dir = os.path.join(out_dir, 'labels')
if not os.path.exists(out_labels_dir):
os.makedirs(out_labels_dir)
data_dirs = {
name: os.path.join(out_dir, f"{name}-dataset")
for name in label_types
}
for name, dir_ in data_dirs.items():
if not os.path.exists(dir_):
os.makedirs(dir_)
event_vocab_df, part_vocab, action_vocab = load_vocabs(
os.path.join(data_dir, 'ANU_ikea_dataset', 'indexing_files',
'atomic_action_list.txt'))
event_vocab_df.to_csv(os.path.join(out_labels_dir, 'event-vocab.csv'))
event_vocab = event_vocab_df.index.tolist()
vocabs = {'event': event_vocab, 'action': action_vocab, 'part': part_vocab}
vocabs = {label_name: vocabs[label_name] for label_name in label_types}
for name, vocab in vocabs.items():
utils.saveVariable(vocab, 'vocab', data_dirs[name])
label_fn = os.path.join(annotation_dir, 'gt_segments.json')
seq_ids, event_labels, metadata = load_action_labels(
label_fn, event_vocab_df)
utils.saveMetadata(metadata, out_labels_dir)
for name, dir_ in data_dirs.items():
utils.saveMetadata(metadata, dir_)
logger.info(f"Loaded {len(seq_ids)} sequences from {label_fn}")
part_names = [name for name in part_vocab if name != '']
col_names = [f"{name}_active" for name in part_names]
integerizers = {
label_name: {name: i
for i, name in enumerate(label_vocab)}
for label_name, label_vocab in vocabs.items()
}
all_slowfast_labels_seg = collections.defaultdict(list)
all_slowfast_labels_win = collections.defaultdict(list)
counts = np.zeros((len(action_vocab), len(part_vocab)), dtype=int)
for i, seq_id in enumerate(seq_ids):
seq_id_str = f"seq={seq_id}"
seq_dir_name = metadata['dir_name'].loc[seq_id]
event_segs = event_labels[i]
if not event_segs.any(axis=None):
logger.warning(f"No event labels for sequence {seq_id}")
continue
event_data = make_event_data(
event_segs,
sorted(glob.glob(os.path.join(frames_dir, seq_dir_name,
'*.jpg'))), integerizers['event'],
integerizers['action'], integerizers['part'],
event_vocab.index('NA'), action_vocab.index('NA'), False)
event_wins = make_window_clips(event_data, vocabs['event'],
vocabs['action'], **win_params)
event_data.to_csv(os.path.join(out_labels_dir,
f"{seq_id_str}_data.csv"),
index=False)
event_segs.to_csv(os.path.join(out_labels_dir,
f"{seq_id_str}_segs.csv"),
index=False)
filenames = event_data['fn'].to_list()
label_indices = {}
for name in label_types:
if name == 'part':
label_indices[name] = event_data[col_names].to_numpy()
seg_labels_slowfast = make_slowfast_labels(
event_segs[['start', 'end']],
getActivePart(event_segs[col_names], part_names),
event_data['fn'],
integerizers[name],
col_format=col_format)
win_labels_slowfast = make_slowfast_labels(
event_wins[['start', 'end']],
getActivePart(event_wins[col_names], part_names),
event_data['fn'],
integerizers[name],
col_format=col_format)
else:
label_indices[name] = event_data[name].to_numpy()
seg_labels_slowfast = make_slowfast_labels(
event_segs[['start', 'end']],
event_segs[name],
event_data['fn'],
integerizers[name],
col_format=col_format)
win_labels_slowfast = make_slowfast_labels(
event_wins[['start', 'end']],
event_wins[name],
event_data['fn'],
integerizers[name],
col_format=col_format)
utils.saveVariable(filenames, f'{seq_id_str}_frame-fns',
data_dirs[name])
utils.saveVariable(label_indices[name], f'{seq_id_str}_labels',
data_dirs[name])
seg_labels_slowfast.to_csv(
os.path.join(data_dirs[name],
f'{seq_id_str}_slowfast-labels.csv'),
**slowfast_csv_params)
win_labels_slowfast.to_csv(
os.path.join(data_dirs[name],
f'{seq_id_str}_slowfast-labels.csv'),
**slowfast_csv_params)
all_slowfast_labels_seg[name].append(seg_labels_slowfast)
all_slowfast_labels_win[name].append(win_labels_slowfast)
plot_event_labels(os.path.join(fig_dir, f"{seq_id_str}.png"),
label_indices['event'], label_indices['action'],
label_indices['part'], event_vocab, action_vocab,
part_names)
for part_activity_row, action_index in zip(label_indices['part'],
label_indices['action']):
for i, is_active in enumerate(part_activity_row):
part_index = integerizers['part'][part_names[i]]
counts[action_index, part_index] += int(is_active)
for name, labels in all_slowfast_labels_seg.items():
pd.concat(labels, axis=0).to_csv(
os.path.join(data_dirs[name], 'slowfast-labels_seg.csv'),
**slowfast_csv_params)
for name, labels in all_slowfast_labels_win.items():
pd.concat(labels, axis=0).to_csv(
os.path.join(data_dirs[name], 'slowfast-labels_win.csv'),
**slowfast_csv_params)
utils.saveVariable(counts, 'action-part-counts', out_labels_dir)
plt.matshow(counts)
plt.xticks(ticks=range(len(part_vocab)),
labels=part_vocab,
rotation='vertical')
plt.yticks(ticks=range(len(action_vocab)), labels=action_vocab)
plt.savefig(os.path.join(fig_dir, 'action-part-coocurrence.png'),
bbox_inches='tight')
def main(out_dir=None,
data_dir=None,
cv_data_dir=None,
score_dirs=[],
fusion_method='sum',
prune_imu=None,
standardize=None,
decode=None,
plot_predictions=None,
results_file=None,
sweep_param_name=None,
gpu_dev_id=None,
model_params={},
cv_params={},
train_params={},
viz_params={}):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
score_dirs = tuple(map(os.path.expanduser, score_dirs))
if cv_data_dir is not None:
cv_data_dir = os.path.expanduser(cv_data_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
def loadAll(seq_ids, var_name, data_dir):
def loadOne(seq_id):
fn = os.path.join(data_dir, f'trial={seq_id}_{var_name}')
return joblib.load(fn)
return tuple(map(loadOne, seq_ids))
device = torchutils.selectDevice(gpu_dev_id)
# Load data
dir_trial_ids = tuple(
set(utils.getUniqueIds(d, prefix='trial=', to_array=True))
for d in score_dirs)
dir_trial_ids += (set(
utils.getUniqueIds(data_dir, prefix='trial=', to_array=True)), )
trial_ids = np.array(list(sorted(set.intersection(*dir_trial_ids))))
for dir_name, t_ids in zip(score_dirs + (data_dir, ), dir_trial_ids):
logger.info(f"{len(t_ids)} trial ids from {dir_name}:")
logger.info(f" {t_ids}")
logger.info(f"{len(trial_ids)} trials in intersection: {trial_ids}")
assembly_seqs = loadAll(trial_ids, 'assembly-seq.pkl', data_dir)
feature_seqs = tuple(
loadAll(trial_ids, 'data-scores.pkl', d) for d in score_dirs)
feature_seqs = tuple(zip(*feature_seqs))
# Combine feature seqs
include_indices = []
for i, seq_feats in enumerate(feature_seqs):
feat_shapes = tuple(f.shape for f in seq_feats)
include_seq = all(f == feat_shapes[0] for f in feat_shapes)
if include_seq:
include_indices.append(i)
else:
warn_str = (
f'Excluding trial {trial_ids[i]} with mismatched feature shapes: '
f'{feat_shapes}')
logger.warning(warn_str)
trial_ids = trial_ids[include_indices]
assembly_seqs = tuple(assembly_seqs[i] for i in include_indices)
feature_seqs = tuple(feature_seqs[i] for i in include_indices)
feature_seqs = tuple(np.stack(f) for f in feature_seqs)
# Define cross-validation folds
if cv_data_dir is None:
dataset_size = len(trial_ids)
cv_folds = utils.makeDataSplits(dataset_size, **cv_params)
cv_fold_trial_ids = tuple(
tuple(map(lambda x: trial_ids[x], splits)) for splits in cv_folds)
else:
fn = os.path.join(cv_data_dir, 'cv-fold-trial-ids.pkl')
cv_fold_trial_ids = joblib.load(fn)
def getSplit(split_idxs):
split_data = tuple(
tuple(s[i] for i in split_idxs)
for s in (feature_seqs, assembly_seqs, trial_ids))
return split_data
gt_scores = []
all_scores = []
num_keyframes_total = 0
num_rgb_errors_total = 0
num_correctable_errors_total = 0
num_oov_total = 0
num_changed_total = 0
for cv_index, (train_ids, test_ids) in enumerate(cv_fold_trial_ids):
try:
test_idxs = np.array(
[trial_ids.tolist().index(i) for i in test_ids])
include_indices.append(cv_index)
except ValueError:
logger.info(
f" Skipping fold {cv_index}: missing test data {test_ids}")
logger.info(f'CV fold {cv_index + 1}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(test_ids)} test)')
# TRAIN PHASE
if cv_data_dir is None:
train_idxs = np.array([trial_ids.index(i) for i in train_ids])
train_assembly_seqs = tuple(assembly_seqs[i] for i in train_idxs)
train_assemblies = []
for seq in train_assembly_seqs:
list(
labels.gen_eq_classes(seq,
train_assemblies,
equivalent=None))
model = None
else:
fn = f'cvfold={cv_index}_train-assemblies.pkl'
train_assemblies = joblib.load(os.path.join(cv_data_dir, fn))
train_idxs = [
i for i in range(len(trial_ids)) if i not in test_idxs
]
fn = f'cvfold={cv_index}_model.pkl'
model = joblib.load(os.path.join(cv_data_dir, fn))
train_features, train_assembly_seqs, train_ids = getSplit(train_idxs)
if False:
train_labels = tuple(
np.array(
list(
labels.gen_eq_classes(assembly_seq,
train_assemblies,
equivalent=None)), )
for assembly_seq in train_assembly_seqs)
train_set = torchutils.SequenceDataset(train_features,
train_labels,
seq_ids=train_ids,
device=device)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=1,
shuffle=True)
train_epoch_log = collections.defaultdict(list)
# val_epoch_log = collections.defaultdict(list)
metric_dict = {
'Avg Loss': metrics.AverageLoss(),
'Accuracy': metrics.Accuracy()
}
criterion = torch.nn.CrossEntropyLoss()
optimizer_ft = torch.optim.Adam(model.parameters(),
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0,
amsgrad=False)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_ft,
step_size=1,
gamma=1.00)
model = FusionClassifier(num_sources=train_features[0].shape[0])
model, last_model_wts = torchutils.trainModel(
model,
criterion,
optimizer_ft,
lr_scheduler,
train_loader, # val_loader,
device=device,
metrics=metric_dict,
train_epoch_log=train_epoch_log,
# val_epoch_log=val_epoch_log,
**train_params)
torchutils.plotEpochLog(train_epoch_log,
subfig_size=(10, 2.5),
title='Training performance',
fn=os.path.join(
fig_dir,
f'cvfold={cv_index}_train-plot.png'))
test_assemblies = train_assemblies.copy()
for feature_seq, gt_assembly_seq, trial_id in zip(
*getSplit(test_idxs)):
gt_seq = np.array(
list(
labels.gen_eq_classes(gt_assembly_seq,
test_assemblies,
equivalent=None)))
if plot_predictions:
assembly_fig_dir = os.path.join(fig_dir, 'assembly-imgs')
if not os.path.exists(assembly_fig_dir):
os.makedirs(assembly_fig_dir)
for i, assembly in enumerate(test_assemblies):
assembly.draw(assembly_fig_dir, i)
# TEST PHASE
accuracies = []
for feature_seq, gt_assembly_seq, trial_id in zip(
*getSplit(test_idxs)):
gt_seq = np.array(
list(
labels.gen_eq_classes(gt_assembly_seq,
test_assemblies,
equivalent=None)))
num_labels = gt_seq.shape[0]
num_features = feature_seq.shape[-1]
if num_labels != num_features:
err_str = (f"Skipping trial {trial_id}: "
f"{num_labels} labels != {num_features} features")
logger.info(err_str)
continue
# Ignore OOV states in ground-truth
sample_idxs = np.arange(feature_seq.shape[-1])
score_idxs = gt_seq[gt_seq < feature_seq.shape[1]]
sample_idxs = sample_idxs[gt_seq < feature_seq.shape[1]]
gt_scores.append(feature_seq[:, score_idxs, sample_idxs])
all_scores.append(feature_seq.reshape(feature_seq.shape[0], -1))
if fusion_method == 'sum':
score_seq = feature_seq.sum(axis=0)
elif fusion_method == 'rgb_only':
score_seq = feature_seq[1]
elif fusion_method == 'imu_only':
score_seq = feature_seq[0]
else:
raise NotImplementedError()
if not decode:
model = None
if model is None:
pred_seq = score_seq.argmax(axis=0)
elif isinstance(model, torch.nn.Module):
inputs = torch.tensor(feature_seq[None, ...],
dtype=torch.float,
device=device)
outputs = model.forward(inputs)
pred_seq = model.predict(outputs)[0].cpu().numpy()
else:
dummy_samples = np.arange(score_seq.shape[1])
pred_seq, _, _, _ = model.viterbi(dummy_samples,
log_likelihoods=score_seq,
ml_decode=(not decode))
pred_assemblies = [train_assemblies[i] for i in pred_seq]
gt_assemblies = [test_assemblies[i] for i in gt_seq]
acc = metrics.accuracy_upto(pred_assemblies,
gt_assemblies,
equivalence=None)
accuracies.append(acc)
rgb_pred_seq = feature_seq[1].argmax(axis=0)
num_changed = np.sum(rgb_pred_seq != pred_seq)
rgb_is_wrong = rgb_pred_seq != gt_seq
num_rgb_errors = np.sum(rgb_is_wrong)
imu_scores = feature_seq[0]
imu_scores_gt = np.array([
imu_scores[s_idx,
t] if s_idx < imu_scores.shape[0] else -np.inf
for t, s_idx in enumerate(gt_seq)
])
imu_scores_rgb = np.array([
imu_scores[s_idx,
t] if s_idx < imu_scores.shape[0] else -np.inf
for t, s_idx in enumerate(rgb_pred_seq)
])
# imu_scores_gt = imu_scores[gt_seq, range(len(rgb_pred_seq))]
best_imu_scores = imu_scores.max(axis=0)
imu_is_right = imu_scores_gt >= best_imu_scores
rgb_pred_score_is_lower = imu_scores_gt > imu_scores_rgb
is_correctable_error = rgb_is_wrong & imu_is_right & rgb_pred_score_is_lower
num_correctable_errors = np.sum(is_correctable_error)
prop_correctable = num_correctable_errors / num_rgb_errors
num_oov = np.sum(gt_seq >= len(train_assemblies))
num_states = len(gt_seq)
num_keyframes_total += num_states
num_rgb_errors_total += num_rgb_errors
num_correctable_errors_total += num_correctable_errors
num_oov_total += num_oov
num_changed_total += num_changed
logger.info(f" trial {trial_id}: {num_states} keyframes")
logger.info(f" accuracy (fused): {acc * 100:.1f}%")
logger.info(
f" {num_oov} OOV states ({num_oov / num_states * 100:.1f}%)"
)
logger.info(
f" {num_rgb_errors} RGB errors; "
f"{num_correctable_errors} correctable from IMU ({prop_correctable * 100:.1f}%)"
)
saveVariable(score_seq, f'trial={trial_id}_data-scores')
saveVariable(pred_assemblies,
f'trial={trial_id}_pred-assembly-seq')
saveVariable(gt_assemblies, f'trial={trial_id}_gt-assembly-seq')
if plot_predictions:
io_figs_dir = os.path.join(fig_dir, 'system-io')
if not os.path.exists(io_figs_dir):
os.makedirs(io_figs_dir)
fn = os.path.join(io_figs_dir, f'trial={trial_id:03}.png')
utils.plot_array(feature_seq, (gt_seq, pred_seq, score_seq),
('gt', 'pred', 'scores'),
fn=fn)
score_figs_dir = os.path.join(fig_dir, 'modality-scores')
if not os.path.exists(score_figs_dir):
os.makedirs(score_figs_dir)
plot_scores(feature_seq,
k=25,
fn=os.path.join(score_figs_dir,
f"trial={trial_id:03}.png"))
paths_dir = os.path.join(fig_dir, 'path-imgs')
if not os.path.exists(paths_dir):
os.makedirs(paths_dir)
assemblystats.drawPath(pred_seq, trial_id,
f"trial={trial_id}_pred-seq", paths_dir,
assembly_fig_dir)
assemblystats.drawPath(gt_seq, trial_id,
f"trial={trial_id}_gt-seq", paths_dir,
assembly_fig_dir)
label_seqs = (gt_seq, ) + tuple(
scores.argmax(axis=0) for scores in feature_seq)
label_seqs = np.row_stack(label_seqs)
k = 10
for i, scores in enumerate(feature_seq):
label_score_seqs = tuple(
np.array([
scores[s_idx,
t] if s_idx < scores.shape[0] else -np.inf
for t, s_idx in enumerate(label_seq)
]) for label_seq in label_seqs)
label_score_seqs = np.row_stack(label_score_seqs)
drawPaths(label_seqs,
f"trial={trial_id}_pred-scores_modality={i}",
paths_dir,
assembly_fig_dir,
path_scores=label_score_seqs)
topk_seq = (-scores).argsort(axis=0)[:k, :]
path_scores = np.column_stack(
tuple(scores[idxs, i]
for i, idxs in enumerate(topk_seq.T)))
drawPaths(topk_seq,
f"trial={trial_id}_topk_modality={i}",
paths_dir,
assembly_fig_dir,
path_scores=path_scores)
label_score_seqs = tuple(
np.array([
score_seq[s_idx,
t] if s_idx < score_seq.shape[0] else -np.inf
for t, s_idx in enumerate(label_seq)
]) for label_seq in label_seqs)
label_score_seqs = np.row_stack(label_score_seqs)
drawPaths(label_seqs,
f"trial={trial_id}_pred-scores_fused",
paths_dir,
assembly_fig_dir,
path_scores=label_score_seqs)
topk_seq = (-score_seq).argsort(axis=0)[:k, :]
path_scores = np.column_stack(
tuple(score_seq[idxs, i]
for i, idxs in enumerate(topk_seq.T)))
drawPaths(topk_seq,
f"trial={trial_id}_topk_fused",
paths_dir,
assembly_fig_dir,
path_scores=path_scores)
if accuracies:
fold_accuracy = float(np.array(accuracies).mean())
# logger.info(f' acc: {fold_accuracy * 100:.1f}%')
metric_dict = {'Accuracy': fold_accuracy}
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
num_unexplained_errors = num_rgb_errors_total - (
num_oov_total + num_correctable_errors_total)
prop_correctable = num_correctable_errors_total / num_rgb_errors_total
prop_oov = num_oov_total / num_rgb_errors_total
prop_unexplained = num_unexplained_errors / num_rgb_errors_total
prop_changed = num_changed_total / num_keyframes_total
logger.info("PERFORMANCE ANALYSIS")
logger.info(
f" {num_rgb_errors_total} / {num_keyframes_total} "
f"RGB errors ({num_rgb_errors_total / num_keyframes_total * 100:.1f}%)"
)
logger.info(f" {num_oov_total} / {num_rgb_errors_total} "
f"RGB errors are OOV ({prop_oov * 100:.1f}%)")
logger.info(f" {num_correctable_errors_total} / {num_rgb_errors_total} "
f"RGB errors are correctable ({prop_correctable * 100:.1f}%)")
logger.info(f" {num_unexplained_errors} / {num_rgb_errors_total} "
f"RGB errors are unexplained ({prop_unexplained * 100:.1f}%)")
logger.info(
f" {num_changed_total} / {num_keyframes_total} "
f"Predictions changed after fusion ({prop_changed * 100:.1f}%)")
gt_scores = np.hstack(tuple(gt_scores))
plot_hists(np.exp(gt_scores),
fn=os.path.join(fig_dir, "score-hists_gt.png"))
all_scores = np.hstack(tuple(all_scores))
plot_hists(np.exp(all_scores),
fn=os.path.join(fig_dir, "score-hists_all.png"))
def main(
out_dir=None, data_dir=None, use_vid_ids_from=None,
output_data=None, magnitude_centering=None, resting_from_gt=None,
remove_before_first_touch=None, include_signals=None, fig_type=None):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
logger.info(f"Reading from: {data_dir}")
logger.info(f"Writing to: {out_dir}")
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def loadAll(seq_ids, var_name, from_dir=data_dir, prefix='trial='):
all_data = tuple(
utils.loadVariable(f"{prefix}{seq_id}_{var_name}", from_dir)
for seq_id in seq_ids
)
return all_data
def saveVariable(var, var_name, to_dir=out_data_dir):
utils.saveVariable(var, var_name, to_dir)
if fig_type is None:
fig_type = 'multi'
# Load data
if use_vid_ids_from is None:
trial_ids = utils.getUniqueIds(data_dir, prefix='trial=', to_array=True)
else:
use_vid_ids_from = os.path.expanduser(use_vid_ids_from)
trial_ids = utils.getUniqueIds(use_vid_ids_from, prefix='trial-', to_array=True)
accel_seqs = loadAll(trial_ids, 'accel-samples.pkl')
gyro_seqs = loadAll(trial_ids, 'gyro-samples.pkl')
action_seqs = loadAll(trial_ids, 'action-seq.pkl')
rgb_timestamp_seqs = loadAll(trial_ids, 'rgb-frame-timestamp-seq.pkl')
def validate_imu(seqs):
def is_valid(d):
return not any(np.isnan(x).any() for x in d.values())
return np.array([is_valid(d) for d in seqs])
imu_is_valid = validate_imu(accel_seqs) & validate_imu(gyro_seqs)
logger.info(
f"Ignoring {(~imu_is_valid).sum()} IMU sequences with NaN-valued samples "
f"(of {len(imu_is_valid)} total)"
)
def chooseValid(seq):
return tuple(x for x, is_valid in zip(seq, imu_is_valid) if is_valid)
trial_ids = np.array(list(chooseValid(trial_ids)))
accel_seqs = chooseValid(accel_seqs)
gyro_seqs = chooseValid(gyro_seqs)
action_seqs = chooseValid(action_seqs)
rgb_timestamp_seqs = chooseValid(rgb_timestamp_seqs)
vocab = []
metadata = utils.loadMetadata(data_dir, rows=trial_ids)
utils.saveMetadata(metadata, out_data_dir)
utils.saveVariable(vocab, 'vocab', out_data_dir)
def norm(x):
norm = np.linalg.norm(imu.getImuSamples(x), axis=1)[:, None]
return norm
accel_mag_seqs = tuple(map(lambda x: dictToArray(x, transform=norm), accel_seqs))
gyro_mag_seqs = tuple(map(lambda x: dictToArray(x, transform=norm), gyro_seqs))
imu_timestamp_seqs = utils.batchProcess(makeTimestamps, accel_seqs, gyro_seqs)
if remove_before_first_touch:
before_first_touch_seqs = utils.batchProcess(
beforeFirstTouch, action_seqs, rgb_timestamp_seqs, imu_timestamp_seqs
)
num_ignored = sum(b is None for b in before_first_touch_seqs)
logger.info(
f"Ignoring {num_ignored} sequences without first-touch annotations "
f"(of {len(before_first_touch_seqs)} total)"
)
trials_missing_first_touch = [
i for b, i in zip(before_first_touch_seqs, trial_ids)
if b is None
]
logger.info(f"Trials without first touch: {trials_missing_first_touch}")
def clip(signal, bool_array):
return signal[~bool_array, ...]
accel_mag_seqs = tuple(
clip(signal, b) for signal, b in zip(accel_mag_seqs, before_first_touch_seqs)
if b is not None
)
gyro_mag_seqs = tuple(
clip(signal, b) for signal, b in zip(gyro_mag_seqs, before_first_touch_seqs)
if b is not None
)
imu_timestamp_seqs = tuple(
clip(signal, b) for signal, b in zip(imu_timestamp_seqs, before_first_touch_seqs)
if b is not None
)
trial_ids = tuple(
x for x, b in zip(trial_ids, before_first_touch_seqs)
if b is not None
)
action_seqs = tuple(
x for x, b in zip(action_seqs, before_first_touch_seqs)
if b is not None
)
rgb_timestamp_seqs = tuple(
x for x, b in zip(rgb_timestamp_seqs, before_first_touch_seqs)
if b is not None
)
assembly_seqs = utils.batchProcess(
parseActions,
action_seqs, rgb_timestamp_seqs, imu_timestamp_seqs
)
if output_data == 'components':
accel_feat_seqs = accel_mag_seqs
gyro_feat_seqs = gyro_mag_seqs
unique_components = {frozenset(): 0}
imu_label_seqs = zip(
*tuple(
labels.componentLabels(*args, unique_components)
for args in zip(action_seqs, rgb_timestamp_seqs, imu_timestamp_seqs)
)
)
saveVariable(unique_components, 'unique_components')
elif output_data == 'pairwise components':
imu_label_seqs = utils.batchProcess(
labels.pairwiseComponentLabels, assembly_seqs,
static_kwargs={'lower_tri_only': True, 'include_action_labels': False}
)
accel_feat_seqs = tuple(map(imu.pairwiseFeats, accel_mag_seqs))
gyro_feat_seqs = tuple(map(imu.pairwiseFeats, gyro_mag_seqs))
else:
raise AssertionError()
signals = {'accel': accel_feat_seqs, 'gyro': gyro_feat_seqs}
if include_signals is None:
include_signals = tuple(signals.keys())
signals = tuple(signals[key] for key in include_signals)
imu_feature_seqs = tuple(np.stack(x, axis=-1).squeeze(axis=-1) for x in zip(*signals))
video_seqs = tuple(zip(imu_feature_seqs, imu_label_seqs, trial_ids))
imu.plot_prediction_eg(video_seqs, fig_dir, fig_type=fig_type, output_data=output_data)
video_seqs = tuple(
zip(assembly_seqs, imu_feature_seqs, imu_timestamp_seqs, imu_label_seqs, trial_ids)
)
for assembly_seq, feature_seq, timestamp_seq, label_seq, trial_id in video_seqs:
id_string = f"trial={trial_id}"
saveVariable(assembly_seq, f'{id_string}_assembly-seq')
saveVariable(feature_seq, f'{id_string}_feature-seq')
saveVariable(timestamp_seq, f'{id_string}_timestamp-seq')
saveVariable(label_seq, f'{id_string}_label-seq')
def main(out_dir=None,
data_dir=None,
event_scores_dir=None,
action_scores_dir=None,
part_scores_dir=None,
as_atomic_events=False,
only_fold=None,
plot_io=None,
prefix='seq=',
results_file=None,
sweep_param_name=None,
model_params={},
cv_params={}):
data_dir = os.path.expanduser(data_dir)
event_scores_dir = os.path.expanduser(event_scores_dir)
action_scores_dir = os.path.expanduser(action_scores_dir)
part_scores_dir = os.path.expanduser(part_scores_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
loader = DataLoader(
{
'event': os.path.join(data_dir, 'event-dataset'),
'action': os.path.join(data_dir, 'action-dataset'),
'part': os.path.join(data_dir, 'part-dataset')
}, {
'event': event_scores_dir,
'action': action_scores_dir,
'part': part_scores_dir
},
prefix=prefix)
logger.info(
f"Loaded ids for {len(loader.seq_ids)} sequences from {data_dir}")
all_metrics = collections.defaultdict(list)
if as_atomic_events:
map_df = pd.read_csv(os.path.join(data_dir, 'labels',
'event-vocab.csv'),
index_col=False,
keep_default_na=False)
else:
map_df = pd.DataFrame(
[[f"{action}({part})", action] +
[part == n for n in loader.vocabs['part'] if n]
for action in loader.vocabs['action']
for part in loader.vocabs['part']],
columns=['event', 'action'] +
[f"{n}_active" for n in loader.vocabs['part'] if n])
event_vocab = map_df['event'].to_list()
utils.saveVariable(event_vocab, 'vocab', out_data_dir)
# Define cross-validation folds
cv_folds = utils.makeDataSplits(len(loader.seq_ids), **cv_params)
utils.saveVariable(cv_folds, 'cv-folds', out_data_dir)
for cv_index, cv_fold in enumerate(cv_folds):
if only_fold is not None and cv_index != only_fold:
continue
train_indices, val_indices, test_indices = cv_fold
logger.info(
f"CV FOLD {cv_index + 1} / {len(cv_folds)}: "
f"{len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} test"
)
action_part_counts = sum(
count_action_part_bigrams(loader.vocabs['action'],
loader.vocabs['part'], la, lp)
for la, lp in loader._getLabels(train_indices,
label_name=['action', 'part'],
load_from_data=True))
model = AttributeClassifier(action_part_counts)
plot_action_part_bigrams(
loader.vocabs['action'], loader.vocabs['part'],
(model.action_part_counts > 0).astype(int),
os.path.join(fig_dir,
f'cvfold={cv_index}_action-part-coocurrence.png'))
for i in test_indices:
seq_id = loader.seq_ids[i]
logger.info(f" Processing sequence {seq_id}...")
action_score_seq, part_score_seq, true_action_seq, true_part_seq = loader[
seq_id]
event_score_seq = model.forward(action_score_seq, part_score_seq)
pred_action_seq, pred_part_seq = model.predict(event_score_seq)
true_event_seq = ap_to_event(true_action_seq, true_part_seq,
map_df, loader.vocabs)
pred_event_seq = ap_to_event(pred_action_seq, pred_part_seq,
map_df, loader.vocabs)
metric_dict = {}
metric_dict = eval_metrics(pred_event_seq,
true_event_seq,
name_prefix='Event ',
append_to=metric_dict)
metric_dict = eval_metrics(pred_action_seq,
true_action_seq,
name_prefix='Action ',
append_to=metric_dict)
metric_dict = eval_metrics(pred_part_seq,
true_part_seq,
name_prefix='Part ',
append_to=metric_dict)
for name, value in metric_dict.items():
logger.info(f" {name}: {value * 100:.2f}%")
all_metrics[name].append(value)
seq_id_str = f"seq={seq_id}"
utils.saveVariable(event_score_seq, f'{seq_id_str}_score-seq',
out_data_dir)
utils.saveVariable(true_event_seq, f'{seq_id_str}_true-label-seq',
out_data_dir)
utils.saveVariable(pred_event_seq, f'{seq_id_str}_pred-label-seq',
out_data_dir)
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
if plot_io:
utils.plot_array(
event_score_seq.reshape(event_score_seq.shape[0], -1).T,
(true_event_seq, pred_event_seq), ('true', 'pred'),
fn=os.path.join(fig_dir, f"seq={seq_id:03d}_event.png"))
utils.plot_array(action_score_seq.T,
(true_action_seq, pred_action_seq),
('true', 'pred'),
fn=os.path.join(
fig_dir, f"seq={seq_id:03d}_action.png"))
utils.plot_array(part_score_seq.T,
(true_part_seq, pred_part_seq),
('true', 'pred'),
fn=os.path.join(fig_dir,
f"seq={seq_id:03d}_part.png"))
def main(out_dir=None,
data_dir=None,
preprocess_dir=None,
segments_dir=None,
keyframe_model_fn=None,
max_seqs=None,
subsample_period=None,
frame_scoring_options={},
frame_selection_options={}):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
preprocess_dir = os.path.expanduser(preprocess_dir)
keyframe_model_fn = os.path.expanduser(keyframe_model_fn)
if segments_dir is not None:
segments_dir = os.path.expanduser(segments_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
def loadFromDir(var_name, dir_name):
return joblib.load(os.path.join(dir_name, f"{var_name}.pkl"))
def loadFromPreprocessDir(var_name):
return joblib.load(os.path.join(preprocess_dir, f"{var_name}.pkl"))
def saveToWorkingDir(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f"{var_name}.pkl"))
trial_ids = utils.getUniqueIds(preprocess_dir,
prefix='trial-',
suffix='.pkl')
keyframe_model = joblib.load(keyframe_model_fn)
models.visualizeKeyframeModel(keyframe_model,
fn=os.path.join(fig_dir,
'keyframe-model.png'))
if max_seqs is not None:
trial_ids = trial_ids[:max_seqs]
for seq_idx, trial_id in enumerate(trial_ids):
logger.info(
f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})"
)
logger.info(f" Loading data...")
trial_str = f"trial-{trial_id}"
rgb_frame_seq = loadFromDir(f'{trial_str}_rgb-frame-seq', data_dir)
segment_frame_seq = loadFromDir(f'{trial_str}_segment-frame-seq',
preprocess_dir)
if segments_dir is not None:
fn = f'trial={trial_id}_pred-segment-seq-rgb'
try:
segments_seq = loadFromDir(fn, segments_dir)
except FileNotFoundError:
logger.info(f" File not found: {fn}")
continue
else:
segments_seq = None
logger.info(f" Scoring frames...")
frame_scores = videoprocessing.scoreFrames(
keyframe_model,
rgb_frame_seq,
segment_frame_seq,
score_kwargs=frame_scoring_options)
segment_keyframe_idxs = videoprocessing.selectSegmentKeyframes(
frame_scores,
segment_labels=segments_seq,
**frame_selection_options)
logger.info(f" Saving output...")
fn = os.path.join(fig_dir, f'{trial_str}_scores-plot.png')
plotScores(frame_scores, segment_keyframe_idxs, fn)
def saveFrames(indices, label):
best_rgb = rgb_frame_seq[indices]
best_seg = segment_frame_seq[indices]
rgb_quantized = np.stack(
tuple(
videoprocessing.quantizeImage(keyframe_model, rgb_img,
segment_img)
for rgb_img, segment_img in zip(best_rgb, best_seg)))
imageprocessing.displayImages(
*best_rgb,
*best_seg,
*rgb_quantized,
num_rows=3,
file_path=os.path.join(fig_dir,
f'{trial_str}_best-frames-{label}.png'))
saveFrames(segment_keyframe_idxs, 'segment')
# Save intermediate results
saveToWorkingDir(frame_scores, f'{trial_str}_frame-scores')
saveToWorkingDir(segment_keyframe_idxs, f'{trial_str}_keyframe-idxs')
def main(out_dir=None,
video_data_dir=None,
imu_data_dir=None,
video_seg_scores_dir=None,
imu_seg_scores_dir=None,
gt_keyframes_dir=None,
label_kwargs={}):
out_dir = os.path.expanduser(out_dir)
video_data_dir = os.path.expanduser(video_data_dir)
imu_data_dir = os.path.expanduser(imu_data_dir)
video_seg_scores_dir = os.path.expanduser(video_seg_scores_dir)
if imu_seg_scores_dir is not None:
imu_seg_scores_dir = os.path.expanduser(imu_seg_scores_dir)
if gt_keyframes_dir is not None:
gt_keyframes_dir = os.path.expanduser(gt_keyframes_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
def loadFromDir(var_name, dir_name):
return joblib.load(os.path.join(dir_name, f"{var_name}.pkl"))
def saveToWorkingDir(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f"{var_name}.pkl"))
trial_ids = utils.getUniqueIds(video_seg_scores_dir,
prefix='trial-',
suffix='.pkl')
all_score_seqs = []
all_action_labels = []
for seq_idx, trial_id in enumerate(trial_ids):
logger.info(
f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})"
)
logger.info(" Loading data...")
score_seq = loadFromDir(f"trial-{trial_id}_frame-scores",
video_seg_scores_dir)
raw_labels = loadFromDir(f"trial-{trial_id}_action-seq",
video_data_dir)
action_labels = makeActivityLabels(raw_labels,
seq_len=score_seq.shape[0],
**label_kwargs)
timestamp_seq = loadFromDir(
f"trial-{trial_id}_rgb-frame-timestamp-seq", video_data_dir)
if timestamp_seq.shape != score_seq.shape:
logger.warning(
f"Video dimensions don't match: "
f"{score_seq.shape} scores, {timestamp_seq.shape} timestamps")
continue
if imu_seg_scores_dir is not None:
try:
imu_score_seq = loadFromDir(f'trial={trial_id}_score-seq',
imu_seg_scores_dir)
imu_score_seq = make_imu_feats(imu_score_seq)
except FileNotFoundError:
logger.info(f" IMU scores not found: trial {trial_id}")
continue
imu_timestamp_seq = loadFromDir(f"trial={trial_id}_timestamp-seq",
imu_data_dir)
if imu_timestamp_seq.shape != imu_score_seq.shape:
logger.warning(
f"IMU dimensions don't match: "
f"{imu_score_seq.shape} scores, {imu_timestamp_seq.shape} timestamps"
)
continue
# Downsample imu scores to match rgb scores
imu_score_seq = utils.resampleSeq(imu_score_seq, imu_timestamp_seq,
timestamp_seq)
imu_timestamp_seq = timestamp_seq
else:
imu_score_seq = None
imu_timestamp_seq = None
logger.info(" Saving output...")
gt_keyframe_fn = os.path.join(gt_keyframes_dir,
f"trial-{trial_id}_gt-keyframe-seq.pkl")
if os.path.exists(gt_keyframe_fn):
gt_keyframes = joblib.load(gt_keyframe_fn)
else:
gt_keyframes = None
trial_str = f"trial={trial_id}"
fn = os.path.join(fig_dir, f'{trial_str}_scores-plot.png')
plotScores(timestamp_seq,
score_seq,
action_labels=action_labels,
raw_labels=raw_labels,
imu_timestamp_seq=imu_timestamp_seq,
imu_score_seq=imu_score_seq,
keyframe_idxs=gt_keyframes,
fn=fn)
all_score_seqs.append(score_seq)
all_action_labels.append(action_labels)
# Save intermediate results
score_seq -= np.nanmean(score_seq)
score_is_nan = np.isnan(score_seq)
score_seq[score_is_nan] = 0
features = (score_seq, score_is_nan.astype(float))
if imu_score_seq is not None:
features += (imu_score_seq, )
feature_seq = np.column_stack(features)
saveToWorkingDir(feature_seq, f'{trial_str}_feature-seq')
saveToWorkingDir(action_labels, f'{trial_str}_label-seq')
all_score_seqs = np.hstack(tuple(all_score_seqs))
all_action_labels = np.hstack(tuple(all_action_labels))
fn = os.path.join(fig_dir, 'score-hists.png')
plotScoreHists(all_score_seqs, all_action_labels, fn=fn)
def main(out_dir=None,
rgb_data_dir=None,
rgb_attributes_dir=None,
rgb_vocab_dir=None,
imu_data_dir=None,
imu_attributes_dir=None,
modalities=['rgb', 'imu'],
gpu_dev_id=None,
plot_predictions=None,
results_file=None,
sweep_param_name=None,
model_params={},
cv_params={},
train_params={},
viz_params={}):
out_dir = os.path.expanduser(out_dir)
rgb_data_dir = os.path.expanduser(rgb_data_dir)
rgb_attributes_dir = os.path.expanduser(rgb_attributes_dir)
rgb_vocab_dir = os.path.expanduser(rgb_vocab_dir)
imu_data_dir = os.path.expanduser(imu_data_dir)
imu_attributes_dir = os.path.expanduser(imu_attributes_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def saveVariable(var, var_name, to_dir=out_data_dir):
utils.saveVariable(var, var_name, to_dir)
# Load data
if modalities == ['rgb']:
trial_ids = utils.getUniqueIds(rgb_data_dir,
prefix='trial=',
to_array=True)
logger.info(f"Processing {len(trial_ids)} videos")
else:
rgb_trial_ids = utils.getUniqueIds(rgb_data_dir,
prefix='trial=',
to_array=True)
imu_trial_ids = utils.getUniqueIds(imu_data_dir,
prefix='trial=',
to_array=True)
trial_ids = np.array(
sorted(set(rgb_trial_ids.tolist()) & set(imu_trial_ids.tolist())))
logger.info(
f"Processing {len(trial_ids)} videos common to "
f"RGB ({len(rgb_trial_ids)} total) and IMU ({len(imu_trial_ids)} total)"
)
device = torchutils.selectDevice(gpu_dev_id)
dataset = FusionDataset(trial_ids,
rgb_attributes_dir,
rgb_data_dir,
imu_attributes_dir,
imu_data_dir,
device=device,
modalities=modalities)
utils.saveMetadata(dataset.metadata, out_data_dir)
saveVariable(dataset.vocab, 'vocab')
# parts_vocab = loadVariable('parts-vocab')
edge_labels = {
'rgb':
utils.loadVariable('part-labels', rgb_vocab_dir),
'imu':
np.stack([
labels.inSameComponent(a, lower_tri_only=True)
for a in dataset.vocab
])
}
# edge_labels = revise_edge_labels(edge_labels, input_seqs)
attribute_labels = tuple(edge_labels[name] for name in modalities)
logger.info('Making transition probs...')
transition_probs = make_transition_scores(dataset.vocab)
saveVariable(transition_probs, 'transition-probs')
model = AttributeModel(*attribute_labels, device=device)
if plot_predictions:
figsize = (12, 3)
fig, axis = plt.subplots(1, figsize=figsize)
axis.imshow(edge_labels['rgb'].T, interpolation='none', aspect='auto')
plt.savefig(os.path.join(fig_dir, "edge-labels.png"))
plt.close()
for i, trial_id in enumerate(trial_ids):
logger.info(f"Processing sequence {trial_id}...")
trial_prefix = f"trial={trial_id}"
true_label_seq = dataset.loadTargets(trial_id)
attribute_feats = dataset.loadInputs(trial_id)
score_seq = model(attribute_feats)
pred_label_seq = model.predict(score_seq)
attribute_feats = attribute_feats.cpu().numpy()
score_seq = score_seq.cpu().numpy()
true_label_seq = true_label_seq.cpu().numpy()
pred_label_seq = pred_label_seq.cpu().numpy()
saveVariable(score_seq.T, f'{trial_prefix}_score-seq')
saveVariable(true_label_seq.T, f'{trial_prefix}_label-seq')
if plot_predictions:
fn = os.path.join(fig_dir, f'{trial_prefix}.png')
utils.plot_array(attribute_feats.T,
(true_label_seq, pred_label_seq, score_seq),
('gt', 'pred', 'scores'),
fn=fn)
metric_dict = eval_metrics(pred_label_seq, true_label_seq)
for name, value in metric_dict.items():
logger.info(f" {name}: {value * 100:.2f}%")
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
def main(
out_dir=None, corpus_name=None, default_annotator=None,
metadata_file=None, metadata_criteria={},
start_from=None, stop_at=None,
start_video_from_first_touch=None, save_video_before_first_touch=None,
subsample_period=None, rgb_as_float=None,
modalities=None, use_annotated_keyframes=None, download_gt_keyframes=None):
if modalities is None:
modalities = ('video', 'imu')
out_dir = os.path.expanduser(out_dir)
metadata_file = os.path.expanduser(metadata_file)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def saveToWorkingDir(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f"{var_name}.pkl"))
selected_frame_indices = slice(None, None, subsample_period)
metadata = loadMetadata(metadata_file, metadata_criteria=metadata_criteria)
metadata.to_csv(os.path.join(out_data_dir, 'metadata.csv'), index=True)
corpus = duplocorpus.DuploCorpus(corpus_name)
trial_ids = metadata['VidID']
annotator_names = {}
for seq_idx, trial_id in enumerate(trial_ids):
if start_from is not None and seq_idx < start_from:
continue
if stop_at is not None and seq_idx > stop_at:
break
logger.info(f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})")
if start_video_from_first_touch:
label_seq = corpus.readLabels(trial_id, default_annotator)[0]
first_touch_idxs = label_seq['start'][label_seq['action'] == 7]
if not len(first_touch_idxs):
logger.info(f" Skipping trial {trial_id}: no first touch annotated")
continue
first_touch_idx = int(first_touch_idxs[0])
# Video is sampled at 30 FPS --> start one second before the first
# touch was annotated (unless that frame would have happened before
# the camera started recording)
start_idx = max(0, first_touch_idx - 30)
selected_frame_indices = slice(start_idx, None, subsample_period)
if use_annotated_keyframes:
annotator = 'Jonathan'
keyframe_idxs = labels.loadKeyframeIndices(corpus, annotator, trial_id)
if not keyframe_idxs.any():
logger.info(f" Skipping trial {trial_id}: no labeled keyframes")
continue
selected_frame_indices = keyframe_idxs
logger.info(" Loading labels...")
action_seq, annotator_name, is_valid = corpus.readLabels(trial_id, default_annotator)
if not is_valid:
logger.info(f" Skipping trial {trial_id}: No labels")
continue
if not (action_seq['action'] == 7).sum():
logger.warning(f' trial {trial_id}: missing first touch labels')
annotator_names[trial_id] = annotator_name
rgb_frame_fn_seq = corpus.getRgbFrameFns(trial_id)
if rgb_frame_fn_seq is None:
logger.info(f" Skipping trial {trial_id}: No RGB frames")
continue
rgb_frame_timestamp_seq = corpus.readRgbTimestamps(trial_id, times_only=True)
if rgb_frame_timestamp_seq is None:
logger.info(f" Skipping trial {trial_id}: No RGB timestamps")
continue
depth_frame_fn_seq = corpus.getDepthFrameFns(trial_id)
if depth_frame_fn_seq is None:
logger.info(f" Skipping trial {trial_id}: No depth frames")
continue
depth_frame_timestamp_seq = corpus.readDepthTimestamps(trial_id, times_only=True)
if depth_frame_timestamp_seq is None:
logger.info(f" Skipping trial {trial_id}: No depth timestamps")
continue
if action_seq['start'].max() >= len(rgb_frame_fn_seq):
logger.info(f" Skipping trial {trial_id}: actions longer than rgb frames")
continue
if action_seq['end'].max() >= len(rgb_frame_fn_seq):
logger.info(f" Skipping trial {trial_id}: actions longer than rgb frames")
continue
action_seq = fixStartEndIndices(
action_seq, rgb_frame_timestamp_seq, selected_frame_indices
)
rgb_frame_timestamp_seq = rgb_frame_timestamp_seq[selected_frame_indices]
depth_frame_timestamp_seq = depth_frame_timestamp_seq[selected_frame_indices]
try:
assembly_seq, is_valid = labels.parseLabelSeq(
action_seq, timestamps=rgb_frame_timestamp_seq,
structure_change_only=True,
)
if not is_valid:
logger.info(f" Skipping trial {trial_id}: Bad labels")
continue
assembly_seq[-1].end_idx = len(rgb_frame_timestamp_seq)
except AssertionError:
logger.info(f" Skipping trial {trial_id}: Bad labels")
continue
trial_str = f"trial={trial_id}"
if 'imu' in modalities:
logger.info(" Loading and saving IMU data...")
accel_seq = imu.loadImuSampleSeq(corpus, trial_id, sensor_name='accel')
gyro_seq = imu.loadImuSampleSeq(corpus, trial_id, sensor_name='gyro')
if not accel_seq:
logger.info(f" Skipping trial {trial_id}: Missing accel features")
continue
if not gyro_seq:
logger.info(f" Skipping trial {trial_id}: Missing gyro features")
continue
imu_bounds = seqBounds(None, None, accel_seq=accel_seq, gyro_seq=gyro_seq)
if imu_bounds is None:
logger.info(f" Skipping trial {trial_id}: No overlapping data samples")
continue
accel_seq = imu.resampleImuSeq(accel_seq, seq_bounds=imu_bounds)
gyro_seq = imu.resampleImuSeq(gyro_seq, seq_bounds=imu_bounds)
saveToWorkingDir(accel_seq, f'{trial_str}_accel-samples')
saveToWorkingDir(gyro_seq, f'{trial_str}_gyro-samples')
if 'video' in modalities:
logger.info(" Loading and saving video data...")
rgb_frame_seq = primesense.loadRgbFrameSeq(
rgb_frame_fn_seq, rgb_frame_timestamp_seq,
stack_frames=True
)
if not len(rgb_frame_seq):
logger.warning(f" Skipping trial {trial_id}: RGB frames are empty!")
continue
if rgb_as_float:
rgb_frame_seq = np.stack(
tuple(skimage.img_as_float(f) for f in rgb_frame_seq),
axis=0
)
saveToWorkingDir(
rgb_frame_seq[selected_frame_indices],
f'{trial_str}_rgb-frame-seq'
)
depth_frame_seq = primesense.loadDepthFrameSeq(
depth_frame_fn_seq, depth_frame_timestamp_seq,
stack_frames=True
)
if not len(rgb_frame_seq):
logger.warning(f" Skipping trial {trial_id}: depth frames are empty!")
continue
saveToWorkingDir(
depth_frame_seq[selected_frame_indices],
f'{trial_str}_depth-frame-seq'
)
if save_video_before_first_touch:
end_idx = max(0, first_touch_idx - 30)
bft_frame_indices = slice(None, end_idx, subsample_period)
saveToWorkingDir(
rgb_frame_seq[bft_frame_indices],
f'{trial_str}_rgb-frame-seq-before-first-touch'
)
saveToWorkingDir(
depth_frame_seq[bft_frame_indices],
f'{trial_str}_depth-frame-seq-before-first-touch'
)
saveToWorkingDir(rgb_frame_fn_seq, f'{trial_str}_rgb-frame-fn-seq')
saveToWorkingDir(rgb_frame_timestamp_seq, f'{trial_str}_rgb-frame-timestamp-seq')
saveToWorkingDir(depth_frame_fn_seq, f'{trial_str}_depth-frame-fn-seq')
saveToWorkingDir(depth_frame_timestamp_seq, f'{trial_str}_depth-frame-timestamp-seq')
saveToWorkingDir(action_seq, f'{trial_str}_action-seq')
saveToWorkingDir(assembly_seq, f'{trial_str}_assembly-seq')
if download_gt_keyframes:
annotator = 'Jonathan'
keyframe_idxs = labels.loadKeyframeIndices(corpus, annotator, trial_id)
if not keyframe_idxs.any():
continue
if start_video_from_first_touch:
raise NotImplementedError()
if subsample_period is not None:
keyframe_idxs = utils.roundToInt(keyframe_idxs / subsample_period)
saveToWorkingDir(keyframe_idxs, f'{trial_str}_gt-keyframe-seq')
trial_ids = sorted(annotator_names.keys())
annotator_names = [annotator_names[t_id] for t_id in trial_ids]
annotator_names = pd.DataFrame({'trial_id': trial_ids, 'annotator_name': annotator_names})
annotator_names.to_csv(os.path.join(out_data_dir, 'annotator_names.csv'), index=False)
def main(out_dir=None,
data_dir=None,
person_masks_dir=None,
bg_masks_dir=None,
sat_thresh=1,
start_from=None,
stop_at=None,
num_disp_imgs=None):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
person_masks_dir = os.path.expanduser(person_masks_dir)
bg_masks_dir = os.path.expanduser(bg_masks_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def loadFromDir(var_name, dir_name):
return joblib.load(os.path.join(dir_name, f"{var_name}.pkl"))
def saveToWorkingDir(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f"{var_name}.pkl"))
trial_ids = utils.getUniqueIds(data_dir, prefix='trial=', to_array=True)
for seq_idx, trial_id in enumerate(trial_ids):
if start_from is not None and seq_idx < start_from:
continue
if stop_at is not None and seq_idx > stop_at:
break
trial_str = f"trial={trial_id}"
logger.info(
f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})"
)
logger.info(" Loading data...")
rgb_frame_seq = loadFromDir(f'{trial_str}_rgb-frame-seq', data_dir)
person_mask_seq = loadFromDir(f'{trial_str}_person-mask-seq',
person_masks_dir)
bg_mask_seq_depth = loadFromDir(f'{trial_str}_bg-mask-seq-depth',
bg_masks_dir)
# bg_mask_seq_rgb = loadFromDir(f'{trial_str}_bg-mask-seq-rgb', bg_masks_dir)
logger.info(" Making segment labels...")
fg_mask_seq = ~bg_mask_seq_depth
seg_labels_seq = np.stack(tuple(
map(makeCoarseSegmentLabels, fg_mask_seq)),
axis=0)
hsv_frame_seq = np.stack(tuple(map(makeHsvFrame, rgb_frame_seq)),
axis=0)
sat_frame_seq = hsv_frame_seq[..., 1]
bg_mask_seq_sat = sat_frame_seq < sat_thresh
seg_labels_seq[person_mask_seq] = 0
seg_labels_seq = np.stack(tuple(
makeFineSegmentLabels(segs, sat)
for segs, sat in zip(seg_labels_seq, bg_mask_seq_sat)),
axis=0)
logger.info(" Saving output...")
saveToWorkingDir(seg_labels_seq.astype(np.uint8),
f'{trial_str}_seg-labels-seq')
plotHsvHist(hsv_frame_seq,
seg_labels_seq,
file_path=os.path.join(fig_dir,
f'{trial_str}_hsv-hists.png'))
if num_disp_imgs is not None:
if rgb_frame_seq.shape[0] > num_disp_imgs:
idxs = np.arange(rgb_frame_seq.shape[0])
np.random.shuffle(idxs)
idxs = idxs[:num_disp_imgs]
else:
idxs = slice(None, None, None)
imageprocessing.displayImages(*(rgb_frame_seq[idxs]),
*(bg_mask_seq_sat[idxs]),
*(bg_mask_seq_depth[idxs]),
*(person_mask_seq[idxs]),
*(seg_labels_seq[idxs]),
num_rows=5,
file_path=os.path.join(
fig_dir,
f'{trial_str}_best-frames.png'))
def main(out_dir=None,
data_dir=None,
assembly_data_dir=None,
scores_dir=None,
event_attr_fn=None,
connection_attr_fn=None,
assembly_attr_fn=None,
only_fold=None,
plot_io=None,
prefix='seq=',
stop_after=None,
background_action='',
model_params={},
cv_params={},
stride=None,
results_file=None,
sweep_param_name=None):
data_dir = os.path.expanduser(data_dir)
assembly_data_dir = os.path.expanduser(assembly_data_dir)
scores_dir = os.path.expanduser(scores_dir)
event_attr_fn = os.path.expanduser(event_attr_fn)
connection_attr_fn = os.path.expanduser(connection_attr_fn)
assembly_attr_fn = os.path.expanduser(assembly_attr_fn)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
misc_dir = os.path.join(out_dir, 'misc')
if not os.path.exists(misc_dir):
os.makedirs(misc_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
seq_ids = utils.getUniqueIds(data_dir,
prefix=prefix,
suffix='labels.*',
to_array=True)
dataset = utils.FeaturelessCvDataset(seq_ids,
data_dir,
prefix=prefix,
label_fn_format='labels')
logger.info(
f"Loaded scores for {len(seq_ids)} sequences from {scores_dir}")
# Define cross-validation folds
cv_folds = utils.makeDataSplits(len(seq_ids), **cv_params)
utils.saveVariable(cv_folds, 'cv-folds', out_data_dir)
# Load event, connection attributes
assembly_vocab = tuple(
tuple(sorted(tuple(sorted(joint)) for joint in a))
for a in utils.loadVariable('vocab', assembly_data_dir))
(*probs, assembly_transition_probs), vocabs = loadPartInfo(
event_attr_fn,
connection_attr_fn,
assembly_attr_fn,
background_action=background_action,
assembly_vocab=assembly_vocab)
event_assembly_scores = event_to_assembly_scores(*probs, vocabs)
assembly_transition_scores = np.log(assembly_transition_probs)
viz_transition_probs(os.path.join(fig_dir, 'action-transitions'),
np.exp(event_assembly_scores), vocabs['event_vocab'])
write_transition_probs(os.path.join(misc_dir, 'action-transitions'),
np.exp(event_assembly_scores),
vocabs['event_vocab'], vocabs['assembly_vocab'])
for cv_index, cv_fold in enumerate(cv_folds):
if only_fold is not None and cv_index != only_fold:
continue
train_indices, val_indices, test_indices = cv_fold
logger.info(
f"CV FOLD {cv_index + 1} / {len(cv_folds)}: "
f"{len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} test"
)
train_data, val_data, test_data = dataset.getFold(cv_fold)
cv_str = f'cvfold={cv_index}'
class_priors, event_dur_probs = count_priors(train_data[0],
len(dataset.vocab),
stride=stride,
approx_upto=0.95,
support_only=True)
event_dur_scores = np.log(event_dur_probs)
event_dur_scores = np.zeros_like(event_dur_scores)
scores = (event_dur_scores, event_assembly_scores,
assembly_transition_scores)
model = decode.AssemblyActionRecognizer(scores, vocabs, model_params)
viz_priors(os.path.join(fig_dir, f'{cv_str}_priors'), class_priors,
event_dur_probs)
model.write_fsts(os.path.join(misc_dir, f'{cv_str}_fsts'))
model.save_vocabs(os.path.join(out_data_dir, f'{cv_str}_model-vocabs'))
for i, (_, seq_id) in enumerate(zip(*test_data)):
if stop_after is not None and i >= stop_after:
break
trial_prefix = f"{prefix}{seq_id}"
if model_params['return_label'] == 'input':
true_seq = utils.loadVariable(f"{trial_prefix}_true-label-seq",
scores_dir)
elif model_params['return_label'] == 'output':
try:
true_seq = utils.loadVariable(f"{trial_prefix}_label-seq",
assembly_data_dir)
true_seq = true_seq[::stride]
except AssertionError:
# logger.info(f' Skipping sequence {seq_id}: {e}')
continue
logger.info(f" Processing sequence {seq_id}...")
event_score_seq = utils.loadVariable(f"{trial_prefix}_score-seq",
scores_dir)
if event_score_seq.shape[0] != true_seq.shape[0]:
err_str = (f'Event scores shape {event_score_seq.shape} '
f'!= labels shape {true_seq.shape}')
raise AssertionError(err_str)
# FIXME: the serialized variables are probs, not log-probs
# event_score_seq = suppress_nonmax(event_score_seq)
# event_score_seq = np.log(event_score_seq)
decode_score_seq = model.forward(event_score_seq)
pred_seq = model.predict(decode_score_seq)
metric_dict = eval_metrics(pred_seq, true_seq)
for name, value in metric_dict.items():
logger.info(f" {name}: {value * 100:.2f}%")
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
utils.saveVariable(decode_score_seq, f'{trial_prefix}_score-seq',
out_data_dir)
utils.saveVariable(pred_seq, f'{trial_prefix}_pred-label-seq',
out_data_dir)
utils.saveVariable(true_seq, f'{trial_prefix}_true-label-seq',
out_data_dir)
if plot_io:
utils.plot_array(event_score_seq.T, (pred_seq.T, true_seq.T),
('pred', 'true'),
fn=os.path.join(fig_dir,
f"seq={seq_id:03d}.png"))
write_labels(
os.path.join(misc_dir, f"seq={seq_id:03d}_pred-seq.txt"),
pred_seq, model.output_vocab.as_raw())
write_labels(
os.path.join(misc_dir, f"seq={seq_id:03d}_true-seq.txt"),
true_seq, model.output_vocab.as_raw())
def main(out_dir=None,
data_dir=None,
model_name=None,
gpu_dev_id=None,
batch_size=None,
learning_rate=None,
model_params={},
cv_params={},
train_params={},
viz_params={},
load_masks_params={},
kornia_tfs={},
only_edge=None,
num_disp_imgs=None,
results_file=None,
sweep_param_name=None):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
io_dir = os.path.join(fig_dir, 'model-io')
if not os.path.exists(io_dir):
os.makedirs(io_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def saveVariable(var, var_name, to_dir=out_data_dir):
return utils.saveVariable(var, var_name, to_dir)
trial_ids = utils.getUniqueIds(data_dir, prefix='trial=', to_array=True)
vocab = [BlockAssembly()
] + [make_single_block_state(i) for i in range(len(defn.blocks))]
for seq_id in trial_ids:
assembly_seq = utils.loadVariable(f"trial={seq_id}_assembly-seq",
data_dir)
for assembly in assembly_seq:
utils.getIndex(assembly, vocab)
parts_vocab, part_labels = labels_lib.make_parts_vocab(
vocab, lower_tri_only=True, append_to_vocab=True)
if only_edge is not None:
part_labels = part_labels[:, only_edge:only_edge + 1]
logger.info(
f"Loaded {len(trial_ids)} sequences; {len(vocab)} unique assemblies")
saveVariable(vocab, 'vocab')
saveVariable(parts_vocab, 'parts-vocab')
saveVariable(part_labels, 'part-labels')
device = torchutils.selectDevice(gpu_dev_id)
if model_name == 'AAE':
Dataset = sim2real.DenoisingDataset
elif model_name == 'Resnet':
Dataset = sim2real.RenderDataset
elif model_name == 'Connections':
Dataset = sim2real.ConnectionDataset
elif model_name == 'Labeled Connections':
Dataset = sim2real.LabeledConnectionDataset
occlusion_masks = loadMasks(**load_masks_params)
if occlusion_masks is not None:
logger.info(f"Loaded {occlusion_masks.shape[0]} occlusion masks")
def make_data(shuffle=True):
dataset = Dataset(
parts_vocab,
part_labels,
vocab,
device=device,
occlusion_masks=occlusion_masks,
kornia_tfs=kornia_tfs,
)
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=shuffle)
return dataset, data_loader
for cv_index, cv_splits in enumerate(range(1)):
cv_str = f"cvfold={cv_index}"
train_set, train_loader = make_data(shuffle=True)
test_set, test_loader = make_data(shuffle=False)
val_set, val_loader = make_data(shuffle=True)
if model_name == 'AAE':
model = sim2real.AugmentedAutoEncoder(train_set.data_shape,
train_set.num_classes)
criterion = torchutils.BootstrappedCriterion(
0.25,
base_criterion=torch.nn.functional.mse_loss,
)
metric_names = ('Reciprocal Loss', )
elif model_name == 'Resnet':
model = sim2real.ImageClassifier(train_set.num_classes,
**model_params)
criterion = torch.nn.CrossEntropyLoss()
metric_names = ('Loss', 'Accuracy')
elif model_name == 'Connections':
model = sim2real.ConnectionClassifier(train_set.label_shape[0],
**model_params)
criterion = torch.nn.BCEWithLogitsLoss()
metric_names = ('Loss', 'Accuracy', 'Precision', 'Recall', 'F1')
elif model_name == 'Labeled Connections':
out_dim = int(part_labels.max()) + 1
num_vertices = len(defn.blocks)
edges = np.column_stack(np.tril_indices(num_vertices, k=-1))
if only_edge is not None:
edges = edges[only_edge:only_edge + 1]
model = sim2real.LabeledConnectionClassifier(
out_dim, num_vertices, edges, **model_params)
if only_edge is not None:
logger.info(f"Class freqs: {train_set.class_freqs}")
# criterion = torch.nn.CrossEntropyLoss(weight=1 / train_set.class_freqs[:, 0])
criterion = torch.nn.CrossEntropyLoss()
else:
criterion = torch.nn.CrossEntropyLoss()
# criterion = torchutils.BootstrappedCriterion(
# 0.25, base_criterion=torch.nn.functional.cross_entropy,
# )
metric_names = ('Loss', 'Accuracy', 'Precision', 'Recall', 'F1')
model = model.to(device=device)
optimizer_ft = torch.optim.Adam(model.parameters(),
lr=learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0,
amsgrad=False)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_ft,
step_size=1,
gamma=1.00)
train_epoch_log = collections.defaultdict(list)
val_epoch_log = collections.defaultdict(list)
metric_dict = {name: metrics.makeMetric(name) for name in metric_names}
model, last_model_wts = torchutils.trainModel(
model,
criterion,
optimizer_ft,
lr_scheduler,
train_loader,
val_loader,
device=device,
metrics=metric_dict,
train_epoch_log=train_epoch_log,
val_epoch_log=val_epoch_log,
**train_params)
# Test model
metric_dict = {name: metrics.makeMetric(name) for name in metric_names}
test_io_batches = torchutils.predictSamples(
model.to(device=device),
test_loader,
criterion=criterion,
device=device,
metrics=metric_dict,
data_labeled=True,
update_model=False,
seq_as_batch=train_params['seq_as_batch'],
return_io_history=True)
metric_str = ' '.join(str(m) for m in metric_dict.values())
logger.info('[TST] ' + metric_str)
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
for pred_seq, score_seq, feat_seq, label_seq, trial_id in test_io_batches:
trial_str = f"trial={trial_id}"
saveVariable(pred_seq.cpu().numpy(), f'{trial_str}_pred-label-seq')
saveVariable(score_seq.cpu().numpy(), f'{trial_str}_score-seq')
saveVariable(label_seq.cpu().numpy(),
f'{trial_str}_true-label-seq')
saveVariable(model, f'{cv_str}_model-best')
if train_epoch_log:
torchutils.plotEpochLog(train_epoch_log,
subfig_size=(10, 2.5),
title='Training performance',
fn=os.path.join(
fig_dir, f'{cv_str}_train-plot.png'))
if val_epoch_log:
torchutils.plotEpochLog(val_epoch_log,
subfig_size=(10, 2.5),
title='Heldout performance',
fn=os.path.join(fig_dir,
f'{cv_str}_val-plot.png'))
if num_disp_imgs is not None:
model.plotBatches(test_io_batches, io_dir, dataset=test_set)
def main(
out_dir=None, data_dir=None, model_name=None,
model_params={}, cv_params={}, train_params={}, viz_params={},
plot_predictions=None, results_file=None, sweep_param_name=None):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, f'results.csv')
write_mode = 'w'
else:
results_file = os.path.expanduser(results_file)
write_mode = 'a'
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
def loadAll(seq_ids, var_name, data_dir):
def loadOne(seq_id):
fn = os.path.join(data_dir, f'trial={seq_id}_{var_name}')
return joblib.load(fn)
return tuple(map(loadOne, seq_ids))
# Load data
trial_ids = utils.getUniqueIds(data_dir, prefix='trial=')
feature_seqs = loadAll(trial_ids, 'feature-seq.pkl', data_dir)
label_seqs = loadAll(trial_ids, 'label-seq.pkl', data_dir)
# Define cross-validation folds
dataset_size = len(trial_ids)
cv_folds = utils.makeDataSplits(dataset_size, **cv_params)
def getSplit(split_idxs):
split_data = tuple(
tuple(s[i] for i in split_idxs)
for s in (feature_seqs, label_seqs, trial_ids)
)
return split_data
for cv_index, cv_splits in enumerate(cv_folds):
train_data, val_data, test_data = tuple(map(getSplit, cv_splits))
train_feats, train_labels, train_ids = train_data
test_feats, test_labels, test_ids = test_data
val_feats, val_labels, val_ids = val_data
logger.info(
f'CV fold {cv_index + 1} / {len(cv_folds)}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(val_ids)} val, {len(test_ids)} test)'
)
input_dim = train_set.num_obsv_dims
output_dim = train_set.num_label_types
elif model_name == 'dummy':
# FIXME
model = None
else:
raise AssertionError()
metric_str = ' '.join(str(m) for m in metric_dict.values())
logger.info('[TST] ' + metric_str)
d = {k: v.value for k, v in metric_dict.items()}
utils.writeResults(
results_file, d, sweep_param_name, model_params,
write_mode=write_mode
)
if plot_predictions:
io_fig_dir = os.path.join(fig_dir, 'model-io')
if not os.path.exists(io_fig_dir):
os.makedirs(io_fig_dir)
label_names = ('gt', 'pred')
preds, scores, inputs, gt_labels, ids = zip(*test_io_history)
for batch in test_io_history:
batch = map(lambda x: x.cpu().numpy(), batch)
for preds, _, inputs, gt_labels, seq_id in zip(*batch):
fn = os.path.join(io_fig_dir, f"trial={seq_id}_model-io.png")
utils.plot_array(inputs, (gt_labels, preds), label_names, fn=fn)
def saveTrialData(pred_seq, score_seq, feat_seq, label_seq, trial_id):
saveVariable(pred_seq, f'trial={trial_id}_pred-label-seq')
saveVariable(score_seq, f'trial={trial_id}_score-seq')
saveVariable(label_seq, f'trial={trial_id}_true-label-seq')
for batch in test_io_history:
batch = map(lambda x: x.cpu().numpy(), batch)
for io in zip(*batch):
saveTrialData(*io)
saveVariable(train_ids, f'cvfold={cv_index}_train-ids')
saveVariable(test_ids, f'cvfold={cv_index}_test-ids')
saveVariable(val_ids, f'cvfold={cv_index}_val-ids')
saveVariable(train_epoch_log, f'cvfold={cv_index}_{model_name}-train-epoch-log')
saveVariable(val_epoch_log, f'cvfold={cv_index}_{model_name}-val-epoch-log')
train_fig_dir = os.path.join(fig_dir, 'train-plots')
if not os.path.exists(train_fig_dir):
os.makedirs(train_fig_dir)
def main(out_dir=None,
data_dir=None,
attr_dir=None,
model_name=None,
gpu_dev_id=None,
batch_size=None,
learning_rate=None,
model_params={},
cv_params={},
train_params={},
viz_params={},
plot_predictions=None,
results_file=None,
sweep_param_name=None):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
attr_dir = os.path.expanduser(attr_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, f'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def loadData(seq_id):
var_name = f"trial-{seq_id}_rgb-frame-seq"
data = joblib.load(os.path.join(data_dir, f'{var_name}.pkl'))
return data.swapaxes(1, 3)
def loadLabels(seq_id):
var_name = f"trial-{seq_id}_label-seq"
return joblib.load(os.path.join(attr_dir, f'{var_name}.pkl'))
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
# Load data
trial_ids = utils.getUniqueIds(data_dir)
label_seqs = tuple(map(loadLabels, trial_ids))
device = torchutils.selectDevice(gpu_dev_id)
# Define cross-validation folds
dataset_size = len(trial_ids)
cv_folds = utils.makeDataSplits(dataset_size, **cv_params)
def getSplit(split_idxs):
split_data = tuple(
tuple(s[i] for i in split_idxs) for s in (label_seqs, trial_ids))
return split_data
for cv_index, cv_splits in enumerate(cv_folds):
train_data, val_data, test_data = tuple(map(getSplit, cv_splits))
criterion = torch.nn.BCEWithLogitsLoss()
labels_dtype = torch.float
train_labels, train_ids = train_data
train_set = torchutils.PickledVideoDataset(loadData,
train_labels,
device=device,
labels_dtype=labels_dtype,
seq_ids=train_ids,
batch_size=batch_size)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=1,
shuffle=True)
test_labels, test_ids = test_data
test_set = torchutils.PickledVideoDataset(loadData,
test_labels,
device=device,
labels_dtype=labels_dtype,
seq_ids=test_ids,
batch_size=batch_size)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=1,
shuffle=False)
val_labels, val_ids = val_data
val_set = torchutils.PickledVideoDataset(loadData,
val_labels,
device=device,
labels_dtype=labels_dtype,
seq_ids=val_ids,
batch_size=batch_size)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=1,
shuffle=True)
logger.info(
f'CV fold {cv_index + 1} / {len(cv_folds)}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(val_ids)} val, {len(test_ids)} test)'
)
if model_name == 'resnet':
# input_dim = train_set.num_obsv_dims
output_dim = train_set.num_label_types
model = ImageClassifier(output_dim,
**model_params).to(device=device)
else:
raise AssertionError()
train_epoch_log = collections.defaultdict(list)
val_epoch_log = collections.defaultdict(list)
metric_dict = {
'Avg Loss': metrics.AverageLoss(),
'Accuracy': metrics.Accuracy(),
'Precision': metrics.Precision(),
'Recall': metrics.Recall(),
'F1': metrics.Fmeasure()
}
optimizer_ft = torch.optim.Adam(model.parameters(),
lr=learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0,
amsgrad=False)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_ft,
step_size=1,
gamma=1.00)
model, last_model_wts = torchutils.trainModel(
model,
criterion,
optimizer_ft,
lr_scheduler,
train_loader,
val_loader,
device=device,
metrics=metric_dict,
train_epoch_log=train_epoch_log,
val_epoch_log=val_epoch_log,
**train_params)
# Test model
metric_dict = {
'Avg Loss': metrics.AverageLoss(),
'Accuracy': metrics.Accuracy(),
'Precision': metrics.Precision(),
'Recall': metrics.Recall(),
'F1': metrics.Fmeasure()
}
test_io_history = torchutils.predictSamples(
model.to(device=device),
test_loader,
criterion=criterion,
device=device,
metrics=metric_dict,
data_labeled=True,
update_model=False,
seq_as_batch=train_params['seq_as_batch'],
return_io_history=True)
metric_str = ' '.join(str(m) for m in metric_dict.values())
logger.info('[TST] ' + metric_str)
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
if plot_predictions:
# imu.plot_prediction_eg(test_io_history, fig_dir, fig_type=fig_type, **viz_params)
imu.plot_prediction_eg(test_io_history, fig_dir, **viz_params)
def saveTrialData(pred_seq, score_seq, feat_seq, label_seq, trial_id):
saveVariable(pred_seq.cpu().numpy(),
f'trial={trial_id}_pred-label-seq')
saveVariable(score_seq.cpu().numpy(),
f'trial={trial_id}_score-seq')
saveVariable(label_seq.cpu().numpy(),
f'trial={trial_id}_true-label-seq')
for io in test_io_history:
saveTrialData(*io)
saveVariable(train_ids, f'cvfold={cv_index}_train-ids')
saveVariable(test_ids, f'cvfold={cv_index}_test-ids')
saveVariable(val_ids, f'cvfold={cv_index}_val-ids')
saveVariable(train_epoch_log,
f'cvfold={cv_index}_{model_name}-train-epoch-log')
saveVariable(val_epoch_log,
f'cvfold={cv_index}_{model_name}-val-epoch-log')
saveVariable(metric_dict,
f'cvfold={cv_index}_{model_name}-metric-dict')
saveVariable(model, f'cvfold={cv_index}_{model_name}-best')
model.load_state_dict(last_model_wts)
saveVariable(model, f'cvfold={cv_index}_{model_name}-last')
torchutils.plotEpochLog(train_epoch_log,
subfig_size=(10, 2.5),
title='Training performance',
fn=os.path.join(
fig_dir,
f'cvfold={cv_index}_train-plot.png'))
if val_epoch_log:
torchutils.plotEpochLog(val_epoch_log,
subfig_size=(10, 2.5),
title='Heldout performance',
fn=os.path.join(
fig_dir,
f'cvfold={cv_index}_val-plot.png'))
def main(out_dir=None,
data_dir=None,
scores_dir=None,
start_from=None,
stop_at=None,
results_file=None,
sweep_param_name=None,
cv_params={}):
data_dir = os.path.expanduser(data_dir)
scores_dir = os.path.expanduser(scores_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
io_dir = os.path.join(fig_dir, 'model-io')
if not os.path.exists(io_dir):
os.makedirs(io_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def loadVariable(var_name, from_dir=scores_dir):
var = utils.loadVariable(var_name, from_dir)
return var
def saveVariable(var, var_name, to_dir=out_data_dir):
utils.saveVariable(var, var_name, out_data_dir)
def makeSeqBatches(unflatten, seq_ids):
d = collections.defaultdict(list)
for batch_index, (seq_index, win_index) in enumerate(unflatten):
seq_id = seq_ids[seq_index]
d[seq_id].append(batch_index)
return d
def loadBatchData(cv_index, batch_index):
prefix = f"cvfold={cv_index}_batch={batch_index}"
batch_score = loadVariable(f"{prefix}_score-seq")
batch_pred = loadVariable(f"{prefix}_pred-label-seq").astype(int)
batch_true = loadVariable(f"{prefix}_true-label-seq").astype(int)
return batch_score, batch_pred, batch_true
vocab = loadVariable('vocab')
parts_vocab = loadVariable('parts-vocab')
edge_labels = loadVariable('part-labels')
saveVariable(vocab, 'vocab')
saveVariable(parts_vocab, 'parts-vocab')
saveVariable(edge_labels, 'part-labels')
trial_ids = utils.getUniqueIds(data_dir,
prefix='trial=',
suffix='',
to_array=True)
cv_folds = utils.makeDataSplits(len(trial_ids), **cv_params)
for cv_index, (__, __, test_indices) in enumerate(cv_folds):
logger.info(f"CV FOLD {cv_index + 1} / {len(cv_folds)}")
test_ids = trial_ids[test_indices]
unflatten = loadVariable(f"cvfold={cv_index}_test-set-unflatten")
flatten = makeSeqBatches(unflatten, test_ids)
for seq_id in test_ids:
logger.info(f" Processing sequence {seq_id}...")
batch_idxs = flatten[seq_id]
score_seq, pred_seq, true_seq = map(
np.vstack,
zip(*tuple(loadBatchData(cv_index, i) for i in batch_idxs)))
trial_prefix = f"trial={seq_id}"
rgb_seq = loadVariable(f"{trial_prefix}_rgb-frame-seq.",
from_dir=data_dir)
if score_seq.shape[0] != rgb_seq.shape[0]:
err_str = f"scores shape {score_seq.shape} != data shape {rgb_seq.shape}"
raise AssertionError(err_str)
saveVariable(score_seq, f"{trial_prefix}_score-seq")
saveVariable(pred_seq, f"{trial_prefix}_pred-label-seq")
saveVariable(true_seq, f"{trial_prefix}_true-label-seq")
def main(
out_dir=None, preds_dir=None, data_dir=None, metric_names=None,
plot_output=None, results_file=None, sweep_param_name=None):
if metric_names is None:
metric_names = ('accuracy', 'edit_score', 'overlap_score')
preds_dir = os.path.expanduser(preds_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
logger.info(f"Writing to: {out_dir}")
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
if os.path.exists(results_file):
os.remove(results_file)
else:
results_file = os.path.expanduser(results_file)
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
def loadAll(seq_ids, var_name, data_dir):
def loadOne(seq_id):
fn = os.path.join(data_dir, f'trial={seq_id}_{var_name}')
return joblib.load(fn)
return tuple(map(loadOne, seq_ids))
assembly_vocab = joblib.load(os.path.join(data_dir, 'assembly-vocab.pkl'))
trial_ids = utils.getUniqueIds(preds_dir, prefix='trial=')
pred_seqs = loadAll(trial_ids, 'pred-label-seq.pkl', preds_dir)
true_seqs = loadAll(trial_ids, 'true-label-seq.pkl', preds_dir)
input_seqs = loadAll(trial_ids, 'score-seq.pkl', preds_dir)
action_vocab = []
action_batch = []
assembly_batch = []
for i, trial_id in enumerate(trial_ids):
logger.info(f"VIDEO {trial_id}:")
pred_assembly_index_seq = pred_seqs[i]
true_assembly_index_seq = true_seqs[i]
pred_action_segs, pred_action_index_seq = actionsFromAssemblies(
pred_assembly_index_seq, assembly_vocab, action_vocab
)
lib_assembly.writeAssemblies(
os.path.join(out_data_dir, f'trial={trial_id}_pred-actions.txt'),
pred_action_segs
)
true_action_segs, true_action_index_seq = actionsFromAssemblies(
true_assembly_index_seq, assembly_vocab, action_vocab
)
lib_assembly.writeAssemblies(
os.path.join(out_data_dir, f'trial={trial_id}_true-actions.txt'),
true_action_segs
)
metric_dict = {}
for name in metric_names:
key = f"{name}_action"
value = getattr(LCTM.metrics, name)(pred_action_index_seq, true_action_index_seq) / 100
metric_dict[key] = value
logger.info(f" {key}: {value * 100:.1f}%")
key = f"{name}_assembly"
value = getattr(LCTM.metrics, name)(
pred_assembly_index_seq, true_assembly_index_seq
) / 100
metric_dict[key] = value
logger.info(f" {key}: {value * 100:.1f}%")
utils.writeResults(results_file, metric_dict, sweep_param_name, {})
input_seq = input_seqs[i]
action_batch.append((pred_action_index_seq, input_seq, true_action_index_seq, trial_id))
assembly_batch.append(
(pred_assembly_index_seq, input_seq, true_assembly_index_seq, trial_id)
)
lib_assembly.writeAssemblies(os.path.join(out_data_dir, 'action-vocab.txt'), action_vocab)
label_names = ('true', 'pred')
for preds, inputs, gt_labels, seq_id in action_batch:
fn = os.path.join(fig_dir, f"trial={seq_id}_model-io_action.png")
utils.plot_array(None, (gt_labels, preds), label_names, fn=fn, labels_together=True)
label_names = ('true', 'pred')
for preds, inputs, gt_labels, seq_id in assembly_batch:
fn = os.path.join(fig_dir, f"trial={seq_id}_model-io_assembly.png")
utils.plot_array(None, (gt_labels, preds), label_names, fn=fn, labels_together=True)
def main(out_dir=None,
data_dir=None,
cv_data_dir=None,
score_dirs=[],
fusion_method='sum',
decode=None,
plot_predictions=None,
results_file=None,
sweep_param_name=None,
gpu_dev_id=None,
model_params={},
cv_params={},
train_params={},
viz_params={}):
if cv_data_dir is None:
raise AssertionError()
def score_features(feature_seq):
if fusion_method == 'sum':
return feature_seq.sum(axis=0)
elif fusion_method == 'rgb_only':
return feature_seq[1]
elif fusion_method == 'imu_only':
return feature_seq[0]
else:
raise NotImplementedError()
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
score_dirs = tuple(map(os.path.expanduser, score_dirs))
if cv_data_dir is not None:
cv_data_dir = os.path.expanduser(cv_data_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
def loadAll(seq_ids, var_name, data_dir):
def loadOne(seq_id):
fn = os.path.join(data_dir, f'trial={seq_id}_{var_name}')
return joblib.load(fn)
return tuple(map(loadOne, seq_ids))
dirs = score_dirs + (data_dir, )
trial_ids = trial_ids_from_dirs(*dirs)
assembly_seqs = loadAll(trial_ids, 'assembly-seq.pkl', data_dir)
feature_seqs = tuple(
loadAll(trial_ids, 'data-scores.pkl', d) for d in score_dirs)
feature_seqs = tuple(zip(*feature_seqs))
# Filter out sequences where feature shape don't match
include_indices = idxs_of_matching_shapes(feature_seqs, trial_ids)
trial_ids = trial_ids[include_indices]
assembly_seqs = tuple(assembly_seqs[i] for i in include_indices)
feature_seqs = tuple(feature_seqs[i] for i in include_indices)
feature_seqs = tuple(np.stack(f) for f in feature_seqs)
# Define cross-validation folds
fn = os.path.join(cv_data_dir, 'cv-fold-trial-ids.pkl')
cv_fold_trial_ids = joblib.load(fn)
def getSplit(split_idxs):
split_data = tuple(
tuple(s[i] for i in split_idxs)
for s in (feature_seqs, assembly_seqs, trial_ids))
return split_data
new_cv_fold_trial_ids = []
for cv_index, (train_ids, test_ids) in enumerate(cv_fold_trial_ids):
new_test_ids = np.array(
[i for i in test_ids if np.any(trial_ids == i)])
new_train_ids = np.array(
[i for i in train_ids if np.any(trial_ids == i)])
new_cv_fold_trial_ids.append((new_train_ids, new_test_ids))
cv_fold_trial_ids = new_cv_fold_trial_ids
for cv_index, (train_ids, test_ids) in enumerate(cv_fold_trial_ids):
if not test_ids.any():
logger.info(f"Skipping CV fold {cv_index + 1}: no test seqs")
continue
logger.info(f'CV fold {cv_index + 1}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(test_ids)} test)')
# TRAIN PHASE
model = joblib.load(
os.path.join(cv_data_dir, f'cvfold={cv_index}_model.pkl'))
train_assemblies = joblib.load(
os.path.join(cv_data_dir,
f'cvfold={cv_index}_train-assemblies.pkl'))
test_idxs = np.array([trial_ids.tolist().index(i) for i in test_ids])
train_idxs = np.array([trial_ids.tolist().index(i) for i in train_ids])
# train_idxs = np.array([i for i in range(len(trial_ids)) if i not in test_idxs])
train_features, train_assembly_seqs, train_ids = getSplit(train_idxs)
test_assemblies = train_assemblies.copy()
for feature_seq, gt_assembly_seq, trial_id in zip(
*getSplit(test_idxs)):
gt_seq = np.array(
list(
labels.gen_eq_classes(gt_assembly_seq,
test_assemblies,
equivalent=None)))
vocab = train_assemblies.copy()
train_gt_seqs = tuple(
np.array(list(labels.gen_eq_classes(seq, vocab, equivalent=None)))
for seq in train_assembly_seqs)
train_vocab_idxs = np.unique(np.hstack(train_gt_seqs))
if np.any(train_vocab_idxs > len(train_assemblies)):
raise AssertionError()
train_assembly_is_oov = np.array([
not np.any(train_vocab_idxs == i)
for i in range(len(train_assemblies))
])
# train_vocab = [
# a for i, a in enumerate(train_assemblies)
# if not train_assembly_is_oov[i]
# ]
# model.states = train_vocab
# tx_probs, _, _, = su.smoothCounts(*su.countSeqs(train_gt_seqs))
# model.psi[~train_assembly_is_oov, ~train_assembly_is_oov] = tx_probs
# model.log_psi[~train_assembly_is_oov, ~train_assembly_is_oov] = np.log(tx_probs)
model.psi[train_assembly_is_oov, train_assembly_is_oov] = 0
model.log_psi[train_assembly_is_oov, train_assembly_is_oov] = -np.inf
# TEST PHASE
accuracies = []
for feature_seq, gt_assembly_seq, trial_id in zip(
*getSplit(test_idxs)):
logger.info(f" testing on trial {trial_id}")
gt_seq = np.array(
list(
labels.gen_eq_classes(gt_assembly_seq,
test_assemblies,
equivalent=None)))
num_labels = gt_seq.shape[0]
num_features = feature_seq.shape[-1]
if num_labels != num_features:
err_str = (f"Skipping trial {trial_id}: "
f"{num_labels} labels != {num_features} features")
logger.info(err_str)
continue
score_seq = score_features(feature_seq)
score_seq[train_assembly_is_oov, :] = -np.inf
if not decode:
# pred_seq = score_seq.argmax(axis=0)
pred_seq = torch.tensor(score_seq).argmax(dim=0).numpy()
else:
dummy_samples = np.arange(score_seq.shape[1])
pred_seq, _, _, _ = model.viterbi(dummy_samples,
log_likelihoods=score_seq,
ml_decode=False)
pred_assemblies = [train_assemblies[i] for i in pred_seq]
gt_assemblies = [test_assemblies[i] for i in gt_seq]
acc = metrics.accuracy_upto(pred_assemblies,
gt_assemblies,
equivalence=None)
accuracies.append(acc)
saveVariable(score_seq, f'trial={trial_id}_data-scores')
saveVariable(pred_assemblies,
f'trial={trial_id}_pred-assembly-seq')
saveVariable(gt_assemblies, f'trial={trial_id}_gt-assembly-seq')
if accuracies:
fold_accuracy = float(np.array(accuracies).mean())
logger.info(f' acc: {fold_accuracy * 100:.1f}%')
metric_dict = {'Accuracy': fold_accuracy}
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
def main(
out_dir=None, data_dir=None, prefix='trial=',
feature_fn_format='feature-seq.pkl', label_fn_format='label_seq.pkl',
slowfast_labels_path=None, cv_params={}, slowfast_csv_params={}):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
# Load data
trial_ids = utils.getUniqueIds(
data_dir, prefix=prefix, suffix=feature_fn_format,
to_array=True
)
dataset = utils.CvDataset(
trial_ids, data_dir,
feature_fn_format=feature_fn_format, label_fn_format=label_fn_format,
vocab=[], prefix=prefix
)
utils.saveMetadata(dataset.metadata, out_data_dir)
# Make folds
dataset_size = len(trial_ids)
cv_folds = utils.makeDataSplits(dataset_size, metadata=dataset.metadata, **cv_params)
save_cv_folds(cv_folds, os.path.join(out_data_dir, 'cv-folds.json'))
split_names = ('train', 'val', 'test')
# Check folds
for cv_index, cv_fold in enumerate(cv_folds):
train_data, val_data, test_data = dataset.getFold(cv_fold)
train_feats, train_labels, train_ids = train_data
test_feats, test_labels, test_ids = test_data
val_feats, val_labels, val_ids = val_data
logger.info(
f'CV fold {cv_index + 1} / {len(cv_folds)}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(val_ids)} val, {len(test_ids)} test)'
)
slowfast_labels_pattern = os.path.join(data_dir, 'slowfast-labels*.csv')
for slowfast_labels_path in glob.glob(slowfast_labels_pattern):
cv_str = f"cvfold={cv_index}"
fn = os.path.basename(slowfast_labels_path)
slowfast_labels = pd.read_csv(
slowfast_labels_path, index_col=0, keep_default_na=False,
**slowfast_csv_params
)
for split_indices, split_name in zip(cv_fold, split_names):
matches = tuple(
slowfast_labels.loc[slowfast_labels['video_name'] == vid_id]
for vid_id in dataset.metadata.iloc[split_indices]['dir_name'].to_list()
)
if matches:
split = pd.concat(matches, axis=0)
split.to_csv(
os.path.join(out_data_dir, f"{cv_str}_{split_name}_{fn}"),
**slowfast_csv_params
)
else:
logger.info(f' Skipping empty slowfast split: {split_name}')
def main(
out_dir=None, data_dir=None, part_symmetries=None,
plot_output=None, results_file=None, sweep_param_name=None, start_from=None):
if part_symmetries is None:
part_symmetries = [
['frontbeam_hole_1', 'frontbeam_hole_2', 'backbeam_hole_1', 'backbeam_hole_2'],
['cushion_hole_1', 'cushion_hole_2']
]
part_symmetries = {
symms[i]: symms
for symms in part_symmetries
for i in range(len(symms))
}
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
debug_dir = os.path.join(out_dir, 'debug')
if not os.path.exists(debug_dir):
os.makedirs(debug_dir)
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
logger.info(f"Reading from: {data_dir}")
logger.info(f"Writing to: {out_dir}")
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
# write_mode = 'w'
else:
results_file = os.path.expanduser(results_file)
# write_mode = 'a'
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
labels_dir = os.path.join(data_dir, 'labels')
with open(os.path.join(labels_dir, 'action_and_part_names.yaml'), 'rt') as f:
names = yaml.safe_load(f)
action_names = names['action_names']
action_name_to_index = {name: i for i, name in enumerate(action_names)}
part_names = names['part_names']
part_name_to_index = {name: i for i, name in enumerate(part_names)}
video_ids = []
all_label_arrs = []
for label_fn in glob.glob(os.path.join(labels_dir, "*.csv")):
video_id = utils.stripExtension(label_fn)
labels_arr = pd.read_csv(label_fn)
all_label_arrs.append(labels_arr)
video_ids.append(video_id)
pose_dir = os.path.join(data_dir, 'poses')
pose_ids = tuple(
video_id
for video_id in video_ids
if os.path.exists(os.path.join(pose_dir, video_id))
)
keep_ids = tuple(v_id in pose_ids for v_id in video_ids)
logger.info(
f"Ignoring {len(keep_ids) - sum(keep_ids)} video(s) with missing data: "
f"{', '.join([v_id for v_id, keep in zip(video_ids, keep_ids) if not keep])}"
)
def filterSeq(seq):
return tuple(x for x, keep_id in zip(seq, keep_ids) if keep_id)
all_label_arrs = filterSeq(all_label_arrs)
video_ids = filterSeq(video_ids)
assembly_vocab = []
label_seqs = []
for i, video_id in enumerate(video_ids):
if start_from is not None and i < start_from:
continue
logger.info("PROCESSING VIDEO {0}: {1}".format(i, video_id))
labels_arr = all_label_arrs[i]
video_dir = os.path.join(pose_dir, video_id)
def loadFile(part_name):
path = os.path.join(video_dir, f'{part_name}.csv')
arr = pd.read_csv(path)
return arr
part_data = tuple(loadFile(part_name) for part_name in part_names)
poses_seq = np.stack(tuple(arr.values for arr in part_data), axis=-1)
feature_seq = relativePose(poses_seq, lower_tri_only=True, magnitude_only=True)
label_seq = actionLabels(
labels_arr, feature_seq.shape[0],
action_name_to_index, part_name_to_index
)
part_pair_names = makePairs(part_names, lower_tri_only=True)
is_possible = possibleConnections(part_pair_names)
feature_seq = feature_seq[:, is_possible, :]
label_seq = label_seq[:, is_possible]
part_pair_names = tuple(n for (b, n) in zip(is_possible, part_pair_names) if b)
utils.plot_multi(
np.moveaxis(feature_seq, (0, 1, 2), (-1, 0, 1)), label_seq.T,
axis_names=part_pair_names, label_name='action',
feature_names=('translation_dist', 'rotation_dist'),
tick_names=[''] + action_names,
fn=os.path.join(fig_dir, f"{video_id}_actions.png")
)
label_seq = assemblyLabels(
labels_arr, feature_seq.shape[0],
assembly_vocab=assembly_vocab, symmetries=part_symmetries
)
# expanded_label_seq = _assemblyLabels(
# labels_arr, feature_seq.shape[0],
# assembly_vocab=assembly_vocab, symmetries=part_symmetries
# )
utils.plot_array(
feature_seq.sum(axis=-1).T, (label_seq,), ('assembly',),
fn=os.path.join(fig_dir, f"{video_id}_assemblies.png")
)
label_seqs.append(label_seq)
label_segments, __ = utils.computeSegments(label_seq)
assembly_segments = [assembly_vocab[i] for i in label_segments]
lib_assembly.writeAssemblies(
os.path.join(debug_dir, f'trial={video_id}_assembly-seq.txt'),
assembly_segments
)
video_id = video_id.replace('_', '-')
saveVariable(feature_seq, f'trial={video_id}_feature-seq')
saveVariable(label_seq, f'trial={video_id}_label-seq')
# saveVariable(expanded_label_seq, f'trial={video_id}_expanded-label-seq')
if False:
from seqtools import utils as su
transition_probs, start_probs, end_probs = su.smoothCounts(
*su.countSeqs(label_seqs)
)
# import pdb; pdb.set_trace()
lib_assembly.writeAssemblies(
os.path.join(debug_dir, 'assembly-vocab.txt'),
assembly_vocab
)
saveVariable(assembly_vocab, 'assembly-vocab')
with open(os.path.join(out_data_dir, 'action-vocab.yaml'), 'wt') as f:
yaml.dump(action_names, f)
with open(os.path.join(out_data_dir, 'part-vocab.yaml'), 'wt') as f:
yaml.dump(part_names, f)
def main(out_dir=None,
predictions_dir=None,
imu_data_dir=None,
video_data_dir=None,
use_gt_segments=None,
model_name=None,
model_params={},
results_file=None,
sweep_param_name=None,
cv_params={},
viz_params={},
plot_predictions=None):
predictions_dir = os.path.expanduser(predictions_dir)
out_dir = os.path.expanduser(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, f'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
def loadFromDir(var_name, dir_name):
return joblib.load(os.path.join(dir_name, f"{var_name}.pkl"))
def saveVariable(var, var_name):
joblib.dump(var, os.path.join(out_data_dir, f'{var_name}.pkl'))
# Load data
trial_ids = utils.getUniqueIds(predictions_dir,
prefix='trial=',
suffix='.pkl')
for seq_idx, trial_id in enumerate(trial_ids):
logger.info(
f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})"
)
logger.info(f" Loading data...")
feature_seq = loadFromDir(f"trial={trial_id}_score-seq",
predictions_dir)
label_seq = loadFromDir(f"trial={trial_id}_true-label-seq",
predictions_dir)
feature_seq = feature_seq[..., 2].swapaxes(0, 1)
label_seq = (label_seq.swapaxes(0, 1) == 2).any(axis=1).astype(int)
trial_str = f"trial={trial_id}"
fn = os.path.join(fig_dir, f'{trial_str}_scores-plot.png')
plotScores(feature_seq.max(axis=1), action_labels=label_seq, fn=fn)
# all_score_seqs.append(score_seq)
# all_action_labels.append(label_seq)
feature_seq -= feature_seq.mean()
saveVariable(feature_seq, f'{trial_str}_feature-seq')
saveVariable(label_seq, f'{trial_str}_label-seq')
def main(out_dir=None,
videos_dir=None,
hand_detections_dir=None,
viz_params={},
plot_output=None,
results_file=None,
sweep_param_name=None):
videos_dir = os.path.expanduser(videos_dir)
hand_detections_dir = os.path.expanduser(hand_detections_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
# write_mode = 'w'
else:
results_file = os.path.expanduser(results_file)
# write_mode = 'a'
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
for i, video_fn in enumerate(glob.glob(os.path.join(videos_dir, '*.avi'))):
video_id = utils.stripExtension(video_fn)
logger.info(f"Processing video {video_id}")
frames = vu.readVideo(video_fn)
hand_detections = airplanecorpus.loadHandDetections(
video_id, dir_name=hand_detections_dir, unflatten=True)
colors = ((0, 255, 0), (0, 0, 255))
for detections, frame in zip(hand_detections, frames):
for i, detection in enumerate(detections):
if np.isnan(detection).any():
continue
vizDetection(frame,
detection,
**viz_params,
bounding_box_color=colors[i])
video_fn = os.path.join(fig_dir, f'{video_id}.avi')
vu.writeVideo(frames, video_fn)
def main(out_dir=None,
marker_pose_dir=None,
marker_bundles_dir=None,
labels_dir=None,
urdf_file=None,
start_from=None,
rename_parts={}):
marker_pose_dir = os.path.expanduser(marker_pose_dir)
marker_bundles_dir = os.path.expanduser(marker_bundles_dir)
labels_dir = os.path.expanduser(labels_dir)
urdf_file = os.path.expanduser(urdf_file)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
out_poses_dir = os.path.join(out_data_dir, 'poses')
if not os.path.exists(out_poses_dir):
os.makedirs(out_poses_dir)
out_labels_dir = os.path.join(out_data_dir, 'labels')
if not os.path.exists(out_labels_dir):
os.makedirs(out_labels_dir)
out_idxs_dir = os.path.join(out_data_dir, 'frame-indices')
if not os.path.exists(out_idxs_dir):
os.makedirs(out_idxs_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
diffs_dir = os.path.join(fig_dir, 'diffs')
if not os.path.exists(diffs_dir):
os.makedirs(diffs_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
urdf_dir, _ = os.path.split(urdf_file)
assembly = Assembly.from_xacro(
urdf_file, params={'$(find furniture_description)/urdf': urdf_dir})
marker_to_part = {}
for bundle_fn in glob.glob(os.path.join(marker_bundles_dir, '*.xml')):
part_name, bundle_data = ikea_utils.getBundleData(bundle_fn)
for marker_index in bundle_data.keys():
marker_to_part[marker_index] = part_name
action_names = frozenset()
part_names = frozenset()
for i, label_fn in enumerate(glob.glob(os.path.join(labels_dir, "*.txt"))):
# Load labels
labels, video_id, camera_name = ikea_utils.readLabels(
label_fn, rename_parts=rename_parts)
if not os.path.exists(os.path.join(marker_pose_dir, video_id)):
logger.warning(f"Skipping video {video_id}: no marker poses")
continue
# Load frame fns
frame_fns = readFrameFns(os.path.join(marker_pose_dir, video_id,
'metadata',
f"frame-fns_{camera_name}.csv"),
names_as_int=True)
logger.info("PROCESSING VIDEO {0}: {1}".format(i, video_id))
out_video_dir = os.path.join(out_poses_dir, video_id)
if not os.path.exists(out_video_dir):
os.makedirs(out_video_dir)
marker_samples = {}
part_to_keys = collections.defaultdict(list)
for file_path in glob.glob(
os.path.join(marker_pose_dir, video_id, '*.csv')):
file_id = ikea_utils.basename(file_path)
camera_name, marker_index = ikea_utils.metadataFromTopicName(
file_id)
part_name = marker_to_part[marker_index]
marker_pose_seq = ikea_utils.readMarkerPoses(file_path,
pose_and_index=True)
if np.isnan(marker_pose_seq[:, 1:]).all():
continue
marker_samples[camera_name, marker_index] = marker_pose_seq
part_to_keys[part_name].append((camera_name, marker_index))
all_part_frame_seqs = []
for part_name, marker_keys in part_to_keys.items():
marker_index_seqs, marker_pose_seqs = collectMarkerPoses(
marker_samples, marker_keys)
part_pose_seq = pose.avgPose(*marker_pose_seqs)
# VISUALIZE POSES
pose.plotPoses(os.path.join(
fig_dir, f"{video_id}_part={part_name}_all-markers.png"),
*marker_pose_seqs,
labels=tuple(', '.join(tuple(map(str, k)))
for k in marker_keys))
pose.plotPoses(
os.path.join(fig_dir, f"{video_id}_part={part_name}.png"),
part_pose_seq)
part_frame_seq = np.column_stack(
tuple(
cameraIndices(marker_index_seqs, marker_keys, name)
for name in CAMERA_NAMES))
all_part_frame_seqs.append(part_frame_seq)
# map part pose to hole poses
hole_pose_seqs = partPosesToHolePoses(part_pose_seq, part_name,
assembly)
hole_pose_seqs = averageSiblingHoles(hole_pose_seqs)
for hole_name, hole_pose_seq in hole_pose_seqs.items():
# SAVE POSES
pd.DataFrame(data=hole_pose_seq,
columns=POSE_VAR_NAMES).to_csv(os.path.join(
out_video_dir, '{0}.csv'.format(hole_name)),
index=False)
all_part_frame_seqs = np.stack(tuple(all_part_frame_seqs), axis=-1)
if not np.all(all_part_frame_seqs == all_part_frame_seqs[:, :, :1]):
# import pdb; pdb.set_trace()
raise AssertionError()
frame_idx_seq = all_part_frame_seqs[..., 0]
# labels = partLabelsToHoleLabels(labels, ignore_sibling_holes=True)
labels = frameNamesToSampleIndices(labels, frame_fns, frame_idx_seq)
# save labels
labels.to_csv(os.path.join(out_labels_dir, f"{video_id}.csv"),
index=False)
pd.DataFrame(data=frame_idx_seq,
columns=CAMERA_NAMES).to_csv(os.path.join(
out_idxs_dir, '{0}.csv'.format(video_id)),
index=False)
action_names = action_names | frozenset(labels['action'].tolist())
part_names = part_names | frozenset(labels['arg1'].tolist())
part_names = part_names | frozenset(labels['arg2'].tolist())
with open(os.path.join(out_labels_dir, 'action_and_part_names.yaml'),
'wt') as f:
yaml.dump(
{
'action_names': list(action_names),
'part_names': list(part_names)
}, f)
def main(out_dir=None,
video_data_dir=None,
keyframe_scores_dir=None,
activity_labels_dir=None,
action_labels_dir=None,
max_seqs=None,
use_gt_activity=False,
use_gt_actions=False,
frame_selection_options={}):
out_dir = os.path.expanduser(out_dir)
video_data_dir = os.path.expanduser(video_data_dir)
keyframe_scores_dir = os.path.expanduser(keyframe_scores_dir)
activity_labels_dir = os.path.expanduser(activity_labels_dir)
action_labels_dir = os.path.expanduser(action_labels_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
out_video_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_video_data_dir):
os.makedirs(out_video_data_dir)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
def loadFromDir(var_name, dir_name):
return joblib.load(os.path.join(dir_name, f"{var_name}.pkl"))
def saveToWorkingDir(var, var_name):
joblib.dump(var, os.path.join(out_video_data_dir, f"{var_name}.pkl"))
trial_ids = utils.getUniqueIds(action_labels_dir,
prefix='trial=',
suffix='.pkl')
if max_seqs is not None:
trial_ids = trial_ids[:max_seqs]
all_costs = []
all_precisions = []
all_recalls = []
for seq_idx, trial_id in enumerate(trial_ids):
logger.info(
f"Processing video {seq_idx + 1} / {len(trial_ids)} (trial {trial_id})"
)
logger.info(f" Loading data...")
trial_str = f"trial-{trial_id}"
rgb_frame_seq = loadFromDir(f'{trial_str}_rgb-frame-seq',
video_data_dir)
frame_scores = loadFromDir(f"{trial_str}_frame-scores",
keyframe_scores_dir)
if use_gt_activity:
fn = f'trial={trial_id}_true-label-seq'
else:
fn = f'trial={trial_id}_pred-label-seq'
activity_label_seq = loadFromDir(fn, activity_labels_dir)
if use_gt_actions:
fn = f'trial={trial_id}_true-label-seq'
else:
fn = f'trial={trial_id}_pred-label-seq'
action_label_seq = loadFromDir(fn, action_labels_dir)
segments_seq = makeSegments(action_label_seq, activity_label_seq)
pred_keyframe_idxs = videoprocessing.selectSegmentKeyframes(
frame_scores,
segment_labels=segments_seq,
**frame_selection_options)
# Measure performance
try:
fn = f'trial-{trial_id}_gt-keyframe-seq'
gt_keyframe_idxs = loadFromDir(fn, video_data_dir)
except FileNotFoundError:
gt_keyframe_idxs = None
# Save and visualize output
logger.info(f" Saving output...")
fn = os.path.join(fig_dir, f'{trial_str}_scores-plot.png')
plotScores(frame_scores,
pred_keyframe_idxs,
fn,
segments_seq=segments_seq,
gt_keyframe_idxs=gt_keyframe_idxs)
def saveFrames(indices, label):
best_rgb = rgb_frame_seq[indices]
imageprocessing.displayImages(
*best_rgb,
num_rows=1,
file_path=os.path.join(fig_dir,
f'{trial_str}_best-frames-{label}.png'))
saveFrames(pred_keyframe_idxs, 'pred')
if gt_keyframe_idxs is not None:
saveFrames(gt_keyframe_idxs, 'gt')
best_costs, false_alarms, misses = evalKeyframes(
pred_keyframe_idxs, gt_keyframe_idxs)
num_fps = len(false_alarms)
num_fns = len(misses)
num_tps = best_costs.shape[0]
precision = num_tps / (num_tps + num_fps)
recall = num_tps / (num_tps + num_fns)
logger.info(f" PRC: {precision * 100:03.1f}%")
logger.info(f" REC: {recall * 100:03.1f}%")
all_costs.append(best_costs)
all_precisions.append(precision)
all_recalls.append(recall)
# Save intermediate results
saveToWorkingDir(pred_keyframe_idxs, f'{trial_str}_keyframe-idxs')
avg_precision = np.array(all_precisions).mean()
avg_recall = np.array(all_recalls).mean()
logger.info(
f"AVG PRC: {avg_precision * 100:03.1f}% ({len(all_precisions)} seqs)")
logger.info(
f"AVG REC: {avg_recall * 100:03.1f}% ({len(all_recalls)} seqs)")
all_costs = np.hstack(all_costs)
fn = os.path.join(fig_dir, "keyframe-hist.png")
plotKeyframeMetrics(all_costs, fn)
def main(out_dir=None, data_dir=None, annotation_dir=None):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
annotation_dir = os.path.expanduser(annotation_dir)
annotation_dir = os.path.join(annotation_dir, 'action_annotations')
vocab_fn = os.path.join(data_dir, 'ANU_ikea_dataset', 'indexing_files',
'atomic_action_list.txt')
with open(vocab_fn, 'rt') as file_:
action_vocab = file_.read().split('\n')
part_names = (label.split('pick up ')[1] for label in action_vocab
if label.startswith('pick up'))
new_action_vocab = tuple(f"{part}" for part in part_names)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_labels_dir = os.path.join(out_dir, 'labels')
if not os.path.exists(out_labels_dir):
os.makedirs(out_labels_dir)
# gt_action = np.load(os.path.join(annotation_dir, 'gt_action.npy'), allow_pickle=True)
with open(os.path.join(annotation_dir, 'gt_segments.json'), 'r') as _file:
gt_segments = json.load(_file)
ann_seqs = {
seq_name: [ann for ann in ann_seq['annotation']]
for seq_name, ann_seq in gt_segments['database'].items()
}
kinem_vocab = [lib_asm.Assembly()]
all_label_index_seqs = collections.defaultdict(list)
for seq_name, ann_seq in ann_seqs.items():
logger.info(f"Processing sequence {seq_name}...")
furn_name, other_name = seq_name.split('/')
goal_state = make_goal_state(furn_name)
label_seq = tuple(ann['label'] for ann in ann_seq)
segment_seq = tuple(ann['segment'] for ann in ann_seq)
start_seq, end_seq = tuple(zip(*segment_seq))
df = pd.DataFrame({
'start': start_seq,
'end': end_seq,
'label': label_seq
})
df = df.loc[df['label'] != 'NA']
if not df.any().any():
warn_str = f"No labels: {furn_name}, {other_name}"
logger.warning(warn_str)
continue
out_path = os.path.join(out_labels_dir, furn_name)
if not os.path.exists(out_path):
os.makedirs(out_path)
try:
new_df = convert_labels(df)
except AssertionError as e:
logger.warning(f" Skipping video: {e}")
continue
new_label_seq = tuple(' '.join(part_name.split()[:-1])
for part_name in new_df['arg1'])
label_index_seq = tuple(
new_action_vocab.index(label) for label in new_label_seq)
all_label_index_seqs[furn_name].append(label_index_seq)
kinem_df = parse_assembly_actions(new_df, kinem_vocab)
kinem_states = tuple(kinem_vocab[i] for i in kinem_df['state'])
if not kinem_states[-1] == goal_state:
warn_str = f" Final structure != goal structure:\n{kinem_states[-1]}"
logger.warning(warn_str)
lib_asm.writeAssemblies(
os.path.join(out_path, f"{other_name}_kinem-state.txt"),
kinem_states)
df.to_csv(os.path.join(out_path, f"{other_name}_human.csv"),
index=False)
new_df.to_csv(os.path.join(out_path, f"{other_name}_kinem-action.csv"),
index=False)
kinem_df.to_csv(os.path.join(out_path,
f"{other_name}_kinem-state.csv"),
index=False)
if not any(label_seq):
logger.warning(f"No labels: {seq_name}")
lib_asm.writeAssemblies(os.path.join(out_labels_dir, "kinem-vocab.txt"),
kinem_vocab)
symbol_table = fstutils.makeSymbolTable(new_action_vocab)
for furn_name, label_index_seqs in all_label_index_seqs.items():
label_fsts = tuple(
fstutils.fromSequence(label_index_seq, symbol_table=symbol_table)
for label_index_seq in label_index_seqs)
union_fst = libfst.determinize(fstutils.easyUnion(*label_fsts))
union_fst.minimize()
# for i, label_fst in enumerate(label_fsts):
# fn = os.path.join(fig_dir, f"{furn_name}-{i}")
# label_fst.draw(
# fn, isymbols=symbol_table, osymbols=symbol_table,
# # vertical=True,
# portrait=True,
# acceptor=True
# )
# gv.render('dot', 'pdf', fn)
fn = os.path.join(fig_dir, f"{furn_name}-union")
union_fst.draw(
fn,
isymbols=symbol_table,
osymbols=symbol_table,
# vertical=True,
portrait=True,
acceptor=True)
gv.render('dot', 'pdf', fn)
def main(out_dir=None,
data_dir=None,
segs_dir=None,
scores_dir=None,
vocab_dir=None,
label_type='edges',
gpu_dev_id=None,
start_from=None,
stop_at=None,
num_disp_imgs=None,
results_file=None,
sweep_param_name=None,
model_params={},
cv_params={}):
data_dir = os.path.expanduser(data_dir)
segs_dir = os.path.expanduser(segs_dir)
scores_dir = os.path.expanduser(scores_dir)
vocab_dir = os.path.expanduser(vocab_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
if results_file is None:
results_file = os.path.join(out_dir, 'results.csv')
else:
results_file = os.path.expanduser(results_file)
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
io_dir_images = os.path.join(fig_dir, 'model-io_images')
if not os.path.exists(io_dir_images):
os.makedirs(io_dir_images)
io_dir_plots = os.path.join(fig_dir, 'model-io_plots')
if not os.path.exists(io_dir_plots):
os.makedirs(io_dir_plots)
out_data_dir = os.path.join(out_dir, 'data')
if not os.path.exists(out_data_dir):
os.makedirs(out_data_dir)
seq_ids = utils.getUniqueIds(scores_dir,
prefix='trial=',
suffix='score-seq.*',
to_array=True)
logger.info(
f"Loaded scores for {len(seq_ids)} sequences from {scores_dir}")
link_vocab = {}
joint_vocab = {}
joint_type_vocab = {}
vocab, parts_vocab, part_labels = load_vocab(link_vocab, joint_vocab,
joint_type_vocab, vocab_dir)
pred_vocab = [] # FIXME
if label_type == 'assembly':
logger.info("Converting assemblies -> edges")
state_pred_seqs = tuple(
utils.loadVariable(f"trial={seq_id}_pred-label-seq", scores_dir)
for seq_id in seq_ids)
state_true_seqs = tuple(
utils.loadVariable(f"trial={seq_id}_true-label-seq", scores_dir)
for seq_id in seq_ids)
edge_pred_seqs = tuple(part_labels[seq] for seq in state_pred_seqs)
edge_true_seqs = tuple(part_labels[seq] for seq in state_true_seqs)
elif label_type == 'edge':
logger.info("Converting edges -> assemblies (will take a few minutes)")
edge_pred_seqs = tuple(
utils.loadVariable(f"trial={seq_id}_pred-label-seq", scores_dir)
for seq_id in seq_ids)
edge_true_seqs = tuple(
utils.loadVariable(f"trial={seq_id}_true-label-seq", scores_dir)
for seq_id in seq_ids)
state_pred_seqs = tuple(
edges_to_assemblies(seq, pred_vocab, parts_vocab, part_labels)
for seq in edge_pred_seqs)
state_true_seqs = tuple(
edges_to_assemblies(seq, vocab, parts_vocab, part_labels)
for seq in edge_true_seqs)
device = torchutils.selectDevice(gpu_dev_id)
dataset = sim2real.LabeledConnectionDataset(
utils.loadVariable('parts-vocab', vocab_dir),
utils.loadVariable('part-labels', vocab_dir),
utils.loadVariable('vocab', vocab_dir),
device=device)
all_metrics = collections.defaultdict(list)
# Define cross-validation folds
cv_folds = utils.makeDataSplits(len(seq_ids), **cv_params)
utils.saveVariable(cv_folds, 'cv-folds', out_data_dir)
for cv_index, cv_fold in enumerate(cv_folds):
train_indices, val_indices, test_indices = cv_fold
logger.info(
f"CV FOLD {cv_index + 1} / {len(cv_folds)}: "
f"{len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} test"
)
train_states = np.hstack(
tuple(state_true_seqs[i] for i in (train_indices)))
train_edges = part_labels[train_states]
# state_train_vocab = np.unique(train_states)
# edge_train_vocab = part_labels[state_train_vocab]
train_freq_bigram, train_freq_unigram = edge_joint_freqs(train_edges)
# state_probs = utils.makeHistogram(len(vocab), train_states, normalize=True)
test_states = np.hstack(
tuple(state_true_seqs[i] for i in (test_indices)))
test_edges = part_labels[test_states]
# state_test_vocab = np.unique(test_states)
# edge_test_vocab = part_labels[state_test_vocab]
test_freq_bigram, test_freq_unigram = edge_joint_freqs(test_edges)
f, axes = plt.subplots(1, 2)
axes[0].matshow(train_freq_bigram)
axes[0].set_title('Train')
axes[1].matshow(test_freq_bigram)
axes[1].set_title('Test')
plt.tight_layout()
plt.savefig(
os.path.join(fig_dir, f"edge-freqs-bigram_cvfold={cv_index}.png"))
f, axis = plt.subplots(1)
axis.stem(train_freq_unigram,
label='Train',
linefmt='C0-',
markerfmt='C0o')
axis.stem(test_freq_unigram,
label='Test',
linefmt='C1--',
markerfmt='C1o')
plt.legend()
plt.tight_layout()
plt.savefig(
os.path.join(fig_dir, f"edge-freqs-unigram_cvfold={cv_index}.png"))
for i in test_indices:
seq_id = seq_ids[i]
logger.info(f" Processing sequence {seq_id}...")
trial_prefix = f"trial={seq_id}"
# I include the '.' to differentiate between 'rgb-frame-seq' and
# 'rgb-frame-seq-before-first-touch'
# rgb_seq = utils.loadVariable(f"{trial_prefix}_rgb-frame-seq.", data_dir)
# seg_seq = utils.loadVariable(f"{trial_prefix}_seg-labels-seq", segs_dir)
score_seq = utils.loadVariable(f"{trial_prefix}_score-seq",
scores_dir)
# if score_seq.shape[0] != rgb_seq.shape[0]:
# err_str = f"scores shape {score_seq.shape} != data shape {rgb_seq.shape}"
# raise AssertionError(err_str)
edge_pred_seq = edge_pred_seqs[i]
edge_true_seq = edge_true_seqs[i]
state_pred_seq = state_pred_seqs[i]
state_true_seq = state_true_seqs[i]
num_types = np.unique(state_pred_seq).shape[0]
num_samples = state_pred_seq.shape[0]
num_total = len(pred_vocab)
logger.info(
f" {num_types} assemblies predicted ({num_total} total); "
f"{num_samples} samples")
# edge_freq_bigram, edge_freq_unigram = edge_joint_freqs(edge_true_seq)
# dist_shift = np.linalg.norm(train_freq_unigram - edge_freq_unigram)
metric_dict = {
# 'State OOV rate': oov_rate_state(state_true_seq, state_train_vocab),
# 'Edge OOV rate': oov_rate_edges(edge_true_seq, edge_train_vocab),
# 'State avg prob, true': state_probs[state_true_seq].mean(),
# 'State avg prob, pred': state_probs[state_pred_seq].mean(),
# 'Edge distribution shift': dist_shift
}
metric_dict = eval_edge_metrics(edge_pred_seq,
edge_true_seq,
append_to=metric_dict)
metric_dict = eval_state_metrics(state_pred_seq,
state_true_seq,
append_to=metric_dict)
for name, value in metric_dict.items():
logger.info(f" {name}: {value * 100:.2f}%")
all_metrics[name].append(value)
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
if num_disp_imgs is not None:
pred_images = tuple(
render(dataset, vocab[seg_label])
for seg_label in utils.computeSegments(state_pred_seq)[0])
imageprocessing.displayImages(
*pred_images,
file_path=os.path.join(
io_dir_images,
f"seq={seq_id:03d}_pred-assemblies.png"),
num_rows=None,
num_cols=5)
true_images = tuple(
render(dataset, vocab[seg_label])
for seg_label in utils.computeSegments(state_true_seq)[0])
imageprocessing.displayImages(
*true_images,
file_path=os.path.join(
io_dir_images,
f"seq={seq_id:03d}_true-assemblies.png"),
num_rows=None,
num_cols=5)
utils.plot_array(score_seq.T,
(edge_true_seq.T, edge_pred_seq.T),
('true', 'pred'),
fn=os.path.join(io_dir_plots,
f"seq={seq_id:03d}.png"))
file_basename = utils.stripExtension(__file__)
config_fn = f"{file_basename}.yaml"
config_file_path = os.path.join(os.path.expanduser('~'), 'repo',
'kinemparse', 'scripts', config_fn)
else:
config_fn = os.path.basename(config_file_path)
if os.path.exists(config_file_path):
with open(config_file_path, 'rt') as config_file:
config = yaml.safe_load(config_file)
else:
config = {}
for k, v in args.items():
if isinstance(v, dict) and k in config:
config[k].update(v)
else:
config[k] = v
# Create output directory, instantiate log file and write config options
out_dir = os.path.expanduser(config['out_dir'])
if not os.path.exists(out_dir):
os.makedirs(out_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
with open(os.path.join(out_dir, config_fn), 'w') as outfile:
yaml.dump(config, outfile)
utils.copyFile(__file__, out_dir)
utils.autoreload_ipython()
main(**config)
def main(out_dir=None, data_dir=None, frames_dir=None, metadata_parent_dir=None, start_from=None):
data_dir = os.path.expanduser(data_dir)
frames_dir = os.path.expanduser(frames_dir)
out_dir = os.path.expanduser(out_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
if metadata_parent_dir is not None:
metadata_parent_dir = os.path.join(metadata_parent_dir)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
pose_dir = os.path.join(data_dir, 'poses')
sample_frame_idx_dir = os.path.join(data_dir, 'frame-indices')
for i, video_dir in enumerate(glob.glob(os.path.join(pose_dir, '*'))):
if start_from is not None and i < start_from:
continue
video_id = os.path.basename(os.path.normpath(video_dir))
logger.info("VISUALIZING VIDEO {0}: {1}".format(i, video_id))
out_video_dir = os.path.join(out_dir, 'images', video_id)
if not os.path.exists(out_video_dir):
os.makedirs(out_video_dir)
if metadata_parent_dir is None:
metadata_dir = os.path.join(video_dir, 'metadata')
else:
metadata_dir = os.path.join(metadata_parent_dir, video_id, 'metadata')
camera_params, bundle_topics, camera_frame_fns = ikea_utils.readMetadata(metadata_dir)
sample_frame_idxs = pd.read_csv(
os.path.join(sample_frame_idx_dir, f"{video_id}.csv")
).astype(int)
camera_params = {
camera_name: cameraParamsToCv2(params)
for camera_name, params in camera_params.items()
}
# First, load all the marker poses into memory
prev_shape = None
camera_marker_poses = {}
for file_path in glob.glob(os.path.join(video_dir, '*.csv')):
part_name = ikea_utils.basename(file_path)
pose_seq = pd.read_csv(file_path).to_numpy()
camera_marker_poses[part_name] = pose_seq
# Make sure all pose sequences have the same shape
if prev_shape is not None and pose_seq.shape != prev_shape:
warn_fmt_str = 'Pose shapes differ: {0} != {1} (truncating to shortest)'
logger.warning(warn_fmt_str.format(pose_seq.shape, prev_shape))
if prev_shape is None or pose_seq.shape[0] < prev_shape[0]:
prev_shape = pose_seq.shape
num_pose_samples = prev_shape[0]
camera_marker_poses = {
key: pose_seq[:num_pose_samples, :]
for key, pose_seq in camera_marker_poses.items()
}
for sample_index in range(num_pose_samples):
frame_idxs = sample_frame_idxs.iloc[sample_index, :]
sample_frame_fns = getFrameFns(frame_idxs, camera_frame_fns)
try:
sample_frames = loadFrames(sample_frame_fns, frames_dir, video_id)
except AssertionError:
continue
marker_samples = {
key: pose_seq[sample_index, :]
for key, pose_seq in camera_marker_poses.items()
}
for part_name, sample in marker_samples.items():
if np.isnan(sample).any():
continue
# Draw each marker's pose to every frame
for camera_name in sample_frames.keys():
intrinsics = camera_params[camera_name]['intrinsics']
extrinsics = camera_params[camera_name]['extrinsics']
marker_position = sample[:3]
marker_orientation = sample[3:]
origin_px, axis_endpoints_px = projectAxes(
marker_position, marker_orientation,
intrinsics, extrinsics
)
sample_frames[camera_name] = drawAxes(
sample_frames[camera_name], origin_px, axis_endpoints_px,
axis_colors=RGB, label=part_name
)
output_frame_fn = "{0:06d}.png".format(sample_index)
output_frame = np.vstack(tuple(sample_frames.values()))
cv2.imwrite(os.path.join(out_video_dir, output_frame_fn), output_frame)
