def convert_to_modes(theta_path):
thetas = np.loadtxt(theta_path)
eigenworms_matrix_path = "EigenWorms.csv"
if not os.path.isfile(eigenworms_matrix_path):
urllib.request.urlretrieve(
"https://raw.githubusercontent.com/iteal/wormpose/master/extras/EigenWorms.csv",
filename="EigenWorms.csv")
eigenworms_matrix = load_eigenworms_matrix(eigenworms_matrix_path)
all_modes = []
for t in thetas:
modes = theta_to_modes(t, eigenworms_matrix)
all_modes.append(modes)
return np.array(all_modes)
def test_theta_to_modes():
eigenworms_matrix = load_eigenworms_matrix(_EIGENWORMS_FILE_PATH)
dims = eigenworms_matrix.shape[1]
input_outputs = [
(np.full(dims, np.nan), np.full(dims, np.nan)),
(np.zeros(dims, dtype=float), np.zeros(dims, dtype=float)),
(np.ones(dims, dtype=float), np.zeros(dims, dtype=float)),
(eigenworms_matrix[:,
0], np.array([1] + [0] * (dims - 1)).astype(float)),
(eigenworms_matrix[:, 1] * 2,
np.array([0, 2] + [0] * (dims - 2)).astype(float)),
]
for input, expected_output in input_outputs:
output = theta_to_modes(input, eigenworms_matrix)
assert np.allclose(output, expected_output, atol=1e-3, equal_nan=True)
def evaluate(dataset_path: str, **kwargs):
"""
Evaluate a trained model by predicting synthetic data and recording the image similarity
:param dataset_path: Root path of the dataset containing videos of worm
"""
args = _parse_arguments(dataset_path, kwargs)
mp.set_start_method("spawn", force=True)
random.seed(args.random_seed)
np.random.seed(args.random_seed)
results_dir = os.path.join(args.experiment_dir, "evaluation")
os.makedirs(results_dir, exist_ok=True)
config = load_config(args.config)
eigenworms_matrix = load_eigenworms_matrix(args.eigenworms_matrix_path)
dataset = load_dataset(
dataset_loader=config.dataset_loader,
dataset_path=dataset_path,
selected_video_names=args.video_names,
resize_options=ResizeOptions(resize_factor=config.resize_factor),
**{WORM_IS_LIGHTER: config.worm_is_lighter},
)
pkl_filenames = _generate_synthetic_data(
dataset,
args.num_process,
args.num_samples,
args.postures_generation,
args.temp_dir,
args.random_seed,
)
keras_model = tf.keras.models.load_model(args.model_path, compile=False)
tf_dataset = tf.data.Dataset.from_generator(
partial(_eval_data_gen, pkl_filenames),
tf.float32,
tf.TensorShape(dataset.image_shape + (1, )),
).batch(args.batch_size)
network_predictions = keras_model.predict(tf_dataset)[:args.num_samples]
shuffled_results = ShuffledResults(random_theta=network_predictions)
ResultsScoring(
frame_preprocessing=dataset.frame_preprocessing,
num_process=args.num_process,
temp_dir=args.temp_dir,
image_shape=dataset.image_shape,
)(
results=shuffled_results,
scoring_data_manager=_ScoringDataManager(pkl_filenames),
)
# Keep the maximum score between the two head/tail options for this evaluation
image_scores = np.max(shuffled_results.scores, axis=1)
# Now calculate the angle error and mode error
angle_error = []
modes_error = []
theta_predictions = []
_, theta_labels = _load_templates(pkl_filenames)
for theta_label, theta_results in zip(theta_labels,
shuffled_results.theta):
dists = [
angle_distance(theta_result, theta_label)
for theta_result in theta_results
]
closest_index = int(np.argmin(dists))
closest_theta = theta_results[closest_index]
theta_predictions.append(closest_theta)
angle_error.append(dists[closest_index])
if eigenworms_matrix is not None:
modes_label = theta_to_modes(theta_label, eigenworms_matrix)
modes_prediction = theta_to_modes(closest_theta, eigenworms_matrix)
mode_error = np.abs(modes_label - modes_prediction)
modes_error.append(mode_error)
np.savetxt(os.path.join(results_dir, "image_score.txt"), image_scores)
np.savetxt(os.path.join(results_dir, "angle_error.txt"), angle_error)
np.savetxt(os.path.join(results_dir, "theta_labels.txt"), theta_labels)
np.savetxt(os.path.join(results_dir, "theta_predictions.txt"),
theta_predictions)
if eigenworms_matrix is not None:
np.savetxt(os.path.join(results_dir, "modes_error.txt"), modes_error)
logger.info(
f"Evaluated model with synthetic data,"
f" average image similarity: {np.mean(image_scores):.4f},"
f" average angle error (degrees): {np.rad2deg(np.mean(angle_error)):.2f}"
)
def convert(theta):
theta_flipped = flip_theta(theta)
modes = theta_to_modes(theta, eigenworms_matrix)[:4]
modes_flipped = theta_to_modes(theta_flipped, eigenworms_matrix)[:4]
return (modes, modes_flipped), (theta, theta_flipped)