def main(out_dir=None,
data_dir=None,
attributes_dir=None,
use_gt_segments=None,
segments_dir=None,
cv_data_dir=None,
ignore_trial_ids=None,
gpu_dev_id=None,
plot_predictions=None,
results_file=None,
sweep_param_name=None,
model_params={},
cv_params={},
train_params={},
viz_params={}):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
attributes_dir = os.path.expanduser(attributes_dir)
if segments_dir is not None:
segments_dir = os.path.expanduser(segments_dir)
if cv_data_dir is not None:
cv_data_dir = os.path.expanduser(cv_data_dir)
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, prefix='trial='):
def loadOne(seq_id):
fn = os.path.join(data_dir, f'{prefix}{seq_id}_{var_name}')
return joblib.load(fn)
return tuple(map(loadOne, seq_ids))
# Load data
trial_ids = utils.getUniqueIds(data_dir, prefix='trial=', to_array=True)
assembly_seqs = loadAll(trial_ids,
'assembly-seq.pkl',
data_dir,
prefix='trial=')
# 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)
test_ids = set().union(*tuple(set(fold[-1]) for fold in cv_fold_trial_ids))
logger.info(f"{len(test_ids)} final test ids: {test_ids}")
def getSplit(split_idxs):
split_data = tuple(
tuple(s[i] for i in split_idxs)
for s in (assembly_seqs, trial_ids))
return split_data
for cv_index, (train_ids, test_ids) in enumerate(cv_fold_trial_ids):
logger.info(f'CV fold {cv_index + 1}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(test_ids)} test)')
try:
test_idxs = torch.tensor(
[trial_ids.tolist().index(i) for i in test_ids])
except ValueError:
logger.info(" Skipping fold: missing test data")
continue
# TRAIN PHASE
if cv_data_dir is None:
train_idxs = torch.tensor([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))
else:
fn = os.path.join(cv_data_dir,
f'cvfold={cv_index}_train-assemblies.pkl')
train_assemblies = joblib.load(fn)
signatures = make_signatures(train_assemblies)
if plot_predictions:
figsize = (12, 3)
fig, axis = plt.subplots(1, figsize=figsize)
axis.imshow(signatures.numpy().T,
interpolation='none',
aspect='auto')
plt.savefig(
os.path.join(fig_dir, f"cvfold={cv_index}_signatures.png"))
plt.close()
# TEST PHASE
accuracies = {'assembly': [], 'assembly_upto_eq': []}
for gt_assembly_seq, trial_id in zip(*getSplit(test_idxs)):
# if ignore_trial_ids is not None and trial_id in ignore_trial_ids:
if False:
logger.info(f" Ignoring trial {trial_id} in test fold")
continue
# FIXME: implement consistent data dimensions during serialization
# (ie samples along rows)
# feature_seq shape is (pairs, samples, classes)
# should be (samples, pairs, classes)
try:
fn = os.path.join(attributes_dir,
f'trial={trial_id}_score-seq.pkl')
feature_seq = joblib.load(fn)
except FileNotFoundError:
logger.info(f" File not found: {fn}")
continue
test_assemblies = train_assemblies.copy()
gt_assembly_id_seq = list(
labels.gen_eq_classes(gt_assembly_seq,
test_assemblies,
equivalent=None))
gt_seq = makeAssemblyLabels(gt_assembly_id_seq, gt_assembly_seq)
if use_gt_segments:
segments = utils.makeSegmentLabels(gt_seq)
elif segments_dir is not None:
var_name = 'segment-seq-imu.pkl'
fn = os.path.join(segments_dir, f'trial={trial_id}_{var_name}')
try:
segments = joblib.load(fn)
except FileNotFoundError:
logger.info(f" File not found: {fn}")
continue
else:
segments = None
feature_seq = torch.tensor(feature_seq, dtype=torch.float)
if segments is not None:
feature_seq = feature_seq.transpose(0, 1)
feature_seq, _ = utils.reduce_over_segments(
feature_seq.numpy(),
segments,
reduction=lambda x: x.mean(axis=0))
feature_seq = torch.tensor(feature_seq.swapaxes(0, 1),
dtype=torch.float)
gt_seq, _ = utils.reduce_over_segments(gt_seq.numpy(),
segments)
gt_seq = torch.tensor(gt_seq, dtype=torch.long)
feature_seq = make_features(feature_seq)
score_seq = score(feature_seq, signatures)
pred_seq = predict(score_seq)
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=components_equivalent)
accuracies['assembly_upto_eq'].append(acc)
acc = metrics.accuracy_upto(pred_assemblies, gt_assemblies)
accuracies['assembly'].append(acc)
# FIXME: Improve data naming convention in decode_keyframes.py
saveVariable(score_seq, f'trial={trial_id}_data-scores')
if plot_predictions:
fn = os.path.join(fig_dir, f'trial={trial_id:03}.png')
utils.plot_array(
feature_seq,
(gt_seq.numpy(), pred_seq.numpy(), score_seq.numpy()),
('gt', 'pred', 'scores'),
fn=fn)
if not any(v for v in accuracies.values()):
continue
metric_dict = {}
for name, fold_accuracies in accuracies.items():
fold_accuracy = float(torch.tensor(fold_accuracies).mean())
logger.info(f' {name}: {fold_accuracy * 100:.1f}%')
metric_dict[name] = fold_accuracy
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)
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, metadata_file=None,
plot_predictions=None, results_file=None, sweep_param_name=None):
logger.info(f"Reading from: {data_dir}")
logger.info(f"Writing to: {out_dir}")
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
if metadata_file is not None:
metadata_file = os.path.expanduser(metadata_file)
metadata = pd.read_csv(metadata_file, index_col=0)
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 loadVariable(var_name):
return joblib.load(os.path.join(data_dir, f'{var_name}.pkl'))
def loadAll(seq_ids, var_name):
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))
trial_ids = utils.getUniqueIds(data_dir, prefix='trial=', to_array=True)
pred_assembly_seqs = loadAll(trial_ids, "pred-assembly-seq.pkl")
gt_assembly_seqs = loadAll(trial_ids, "gt-assembly-seq.pkl")
all_assemblies = []
gt_label_seqs = tuple(
np.array(list(labels.gen_eq_classes(gt_assembly_seq, all_assemblies)))
for gt_assembly_seq in gt_assembly_seqs
)
pred_label_seqs = tuple(
np.array(list(labels.gen_eq_classes(pred_assembly_seq, all_assemblies)))
for pred_assembly_seq in pred_assembly_seqs
)
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(all_assemblies):
assembly.draw(assembly_fig_dir, i)
logger.info(f"Evaluating {len(trial_ids)} sequence predictions")
accuracies = {
'state': [],
'is_error': [],
'is_layer': []
}
data = zip(trial_ids, pred_assembly_seqs, gt_assembly_seqs)
for i, trial_id in enumerate(trial_ids):
pred_assembly_seq = pred_assembly_seqs[i]
gt_assembly_seq = gt_assembly_seqs[i]
task = int(metadata.iloc[trial_id]['task id'])
goal = labels.constructGoalState(task)
is_error = functools.partial(is_goal_error, goal)
is_layer = functools.partial(is_goal_layer, goal)
state = None
logger.info(f"SEQUENCE {trial_id}: {len(gt_assembly_seq)} items")
for name in accuracies.keys():
if name == 'is_layer':
pred_is_layer = np.array(list(map(is_layer, pred_assembly_seq)))
gt_is_layer = np.array(list(map(is_layer, gt_assembly_seq)))
matches = pred_is_layer == gt_is_layer
acc = matches.sum() / len(matches)
logger.info(f" {name}: {acc:.2}")
logger.info(f" {gt_is_layer.sum():2} gt layers")
logger.info(f" {pred_is_layer.sum():2} pred layers")
else:
acc = metrics.accuracy_upto(
# pred_assembly_seq[1:], gt_assembly_seq[1:],
pred_assembly_seq, gt_assembly_seq,
equivalence=locals()[name]
)
logger.info(f" {name}: {acc:.2}")
accuracies[name].append(acc)
if plot_predictions:
paths_dir = os.path.join(fig_dir, 'path-imgs')
if not os.path.exists(paths_dir):
os.makedirs(paths_dir)
fn = f"trial={trial_id}_paths"
paths = np.row_stack((gt_label_seqs[i], pred_label_seqs[i]))
path_labels = np.row_stack((gt_is_layer, pred_is_layer))
drawPaths(paths, fn, paths_dir, assembly_fig_dir, path_labels=path_labels)
logger.info("EVALUATION RESULTS:")
max_width = max(map(len, accuracies.keys()))
for name, vals in accuracies.items():
vals = np.array(vals) * 100
mean = vals.mean()
std = vals.std()
logger.info(f" {name:{max_width}}: {mean:4.1f} +/- {std:4.1f}%")
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,
cv_data_dir=None,
scores_dir=None,
eq_class='state index',
plot_predictions=None,
results_file=None,
sweep_param_name=None,
model_params={},
cv_params={},
train_params={},
viz_params={}):
data_dir = os.path.expanduser(data_dir)
out_dir = os.path.expanduser(out_dir)
scores_dir = os.path.expanduser(scores_dir)
if cv_data_dir is not None:
cv_data_dir = os.path.expanduser(cv_data_dir)
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 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(scores_dir, prefix='trial=')
assembly_seqs = loadAll(trial_ids, 'assembly-seq.pkl', data_dir)
score_seqs = loadAll(trial_ids, 'data-scores.pkl', scores_dir)
oov_rate, contextual_oov_rate, state_counts = OOVrate(assembly_seqs,
eq_class=eq_class)
logger.info(f"OOV RATE: {oov_rate * 100:.1f}%")
plotOOV(contextual_oov_rate,
state_counts,
fn=os.path.join(fig_dir, "oovs.png"),
eq_class=eq_class)
scatterOOV(contextual_oov_rate,
state_counts,
fn=os.path.join(fig_dir, "oovs_scatter.png"),
eq_class=eq_class)
import sys
sys.exit()
# 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, f'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 (score_seqs, assembly_seqs, trial_ids))
return split_data
for cv_index, (train_ids, test_ids) in enumerate(cv_fold_trial_ids):
logger.info(f'CV fold {cv_index + 1}: {len(trial_ids)} total '
f'({len(train_ids)} train, {len(test_ids)} test)')
try:
test_idxs = np.array(
[trial_ids.tolist().index(i) for i in test_ids])
except ValueError:
logger.info(f" Skipping fold: missing test data")
continue
# 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))
model = None
train_features, train_assembly_seqs, train_ids = getSplit(train_idxs)
test_assemblies = train_assemblies.copy()
for score_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)))
# oov_rate = OOVrate(train_assembly_seqs)
# logger.info(f" OOV RATE: {oov_rate * 100:.1f}%")
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 score_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_scores = score_seq.shape[-1]
if num_labels != num_scores:
err_str = f"Skipping trial {trial_id}: {num_labels} labels != {num_scores} scores"
logger.info(err_str)
continue
if model is None:
pred_seq = score_seq.argmax(axis=0)
else:
raise AssertionError()
pred_seq = score_seq.argmax(axis=0)
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)
# num_states = len(gt_seq)
# logger.info(f" trial {trial_id}: {num_states} keyframes")
# logger.info(f" accuracy (fused): {acc * 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 accuracies:
fold_accuracy = float(np.array(accuracies).mean())
metric_dict = {'Accuracy': fold_accuracy}
utils.writeResults(results_file, metric_dict, sweep_param_name,
model_params)