def test_reload_eigenworms_file(tmp_path):
eigenworms_matrix = load_eigenworms_matrix(_EIGENWORMS_FILE_PATH)
new_file_path = tmp_path / "eigenworms.csv"
np.savetxt(new_file_path, eigenworms_matrix, delimiter=",")
eigenworms_matrix_reloaded = load_eigenworms_matrix(new_file_path)
assert type(eigenworms_matrix_reloaded) == type(eigenworms_matrix)
assert eigenworms_matrix_reloaded.shape == eigenworms_matrix.shape
assert eigenworms_matrix_reloaded.dtype == eigenworms_matrix.dtype
assert np.allclose(eigenworms_matrix_reloaded, eigenworms_matrix)
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 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_modes_to_theta():
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.array([1] + [0] * (dims - 1)).astype(float), eigenworms_matrix[:,
0]),
(np.array([0] * (dims - 1) + [-10]).astype(float),
eigenworms_matrix[:, -1] * -10),
(np.array([3] + [0] * (dims - 1)).astype(float),
eigenworms_matrix[:, 0] * 3),
]
for input, expected_output in input_outputs:
output = modes_to_theta(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 post_process(dataset_path: str, **kwargs):
"""
Process the raw network results with interpolation and smoothing
:param dataset_path: Root path of the dataset containing videos of worm
"""
args = _parse_arguments(dataset_path, kwargs)
results_root_dir = os.path.join(args.experiment_dir,
default_paths.RESULTS_DIR)
eigenworms_matrix = load_eigenworms_matrix(args.eigenworms_matrix_path)
config = load_config(args.config)
dataset = load_dataset(config.dataset_loader, dataset_path)
spline_interpolation = _SplineInterpolation()
results_files = list(
sorted(glob.glob(os.path.join(results_root_dir, "*",
RESULTS_FILENAME))))
if len(results_files) == 0:
raise FileNotFoundError("No results file to analyze was found")
for results_file in results_files:
video_name = os.path.basename(os.path.dirname(results_file))
with h5py.File(results_file, "r") as results_f:
try:
results_raw = BaseResults(
theta=results_f["resolved"]["theta"][:],
skeletons=results_f["resolved"]["skeletons"][:],
scores=results_f["resolved"]["scores"][:],
)
except Exception:
logger.error(f"Couldn't read results in file {results_file}.")
continue
results_orig = OriginalResults(
theta=results_f["original"]["theta"][:],
skeletons=results_f["original"]["skeletons"][:])
features = dataset.features_dataset[video_name]
missing_values = np.any(np.isnan(results_raw.theta), axis=1)
if missing_values.sum() == len(results_raw.theta):
logger.warning(
f"No valid result was found, stopping postprocessing for {video_name}"
)
continue
segments_boundaries = _get_valid_segments(
is_valid_series=~missing_values,
max_gap_size=args.max_gap_size,
min_segment_size=args.min_segment_size,
)
# interpolate and smooth in angles space
thetas_interp = spline_interpolation.interpolate_tseries(
results_raw.theta, segments_boundaries, args.std_fraction)
results_interp = _calculate_skeleton(thetas_interp, args, dataset,
video_name)
thetas_smooth = _smooth_tseries(
thetas_interp,
args.smoothing_window,
args.poly_order,
segments_boundaries,
)
results_smooth = _calculate_skeleton(thetas_smooth, args, dataset,
video_name)
flipped = False
if features.ventral_side == "clockwise":
results_orig.theta = _dorsal_ventral_flip_theta(
results_orig.theta)
results_raw.theta = _dorsal_ventral_flip_theta(
results_raw.theta)
results_interp.theta = _dorsal_ventral_flip_theta(
results_interp.theta)
results_smooth.theta = _dorsal_ventral_flip_theta(
results_smooth.theta)
flipped = True
if eigenworms_matrix is not None:
setattr(
results_orig, "modes",
_thetas_to_modes(results_orig.theta, eigenworms_matrix))
setattr(results_raw, "modes",
_thetas_to_modes(results_raw.theta, eigenworms_matrix))
setattr(
results_interp, "modes",
_thetas_to_modes(results_interp.theta, eigenworms_matrix))
setattr(
results_smooth, "modes",
_thetas_to_modes(results_smooth.theta, eigenworms_matrix))
# save results
results_saver = ResultsSaver(
temp_dir=args.temp_dir,
results_root_dir=results_root_dir,
results_filename=POSTPROCESSED_RESULTS_FILENAME)
metadata = {
"max_gap_size": args.max_gap_size,
"min_segment_size": args.min_segment_size,
"smoothing_window": args.smoothing_window,
"poly_order": args.poly_order,
"std_fraction": args.std_fraction,
"dorsal_ventral_flip": flipped,
}
results_saver.save(
results={
"orig": results_orig,
"raw": results_raw,
"interp": results_interp,
"smooth": results_smooth
},
metadata=metadata,
video_name=video_name,
)
logger.info(
f"Post-processed worm: {video_name} {'(flipped dorsal-ventral)' if flipped else ''}"
)
def test_load_eigenworms_file():
eigenworms_matrix = load_eigenworms_matrix(_EIGENWORMS_FILE_PATH)
assert type(eigenworms_matrix) is np.ndarray
assert eigenworms_matrix.shape[0] == eigenworms_matrix.shape[1]
assert eigenworms_matrix.dtype == float
def test_load_invalid_path():
with pytest.raises(OSError):
load_eigenworms_matrix("invalid_path")
def test_load_none():
eigenworms_matrix = load_eigenworms_matrix(None)
assert eigenworms_matrix is None
def export_as_images(dataset_loader: str,
dataset_path: str,
video_name: str,
results_path: str,
eigenworms_matrix_path: str,
out_dir: str,
num_process,
temp_dir,
image_size=128):
if out_dir is None:
out_dir = "out"
if num_process is None:
num_process = os.cpu_count()
if temp_dir is None:
temp_dir = tempfile.gettempdir()
temp_dir = tempfile.mkdtemp(dir=temp_dir)
out_dir = os.path.join(out_dir, video_name)
if os.path.exists(out_dir):
shutil.rmtree(out_dir)
os.makedirs(out_dir, exist_ok=True)
eigenworms_matrix = load_eigenworms_matrix(eigenworms_matrix_path)
dataset = load_dataset(dataset_loader=dataset_loader,
dataset_path=dataset_path,
selected_video_names=[video_name],
resize_options=ResizeOptions(image_size=image_size))
features = dataset.features_dataset[video_name]
scoring_data_manager = ScoringDataManager(
video_name=video_name,
frames_dataset=dataset.frames_dataset,
features=features,
)
results_scoring = ResultsScoring(
frame_preprocessing=dataset.frame_preprocessing,
num_process=num_process,
temp_dir=temp_dir,
image_shape=(image_size, image_size),
)
results = np.loadtxt(results_path)
modes = results[:, :5]
theta_mean = results[:, 5]
# convert RCS results to theta
thetas = []
for m, t_m in zip(modes, theta_mean):
theta = modes_to_theta(m, eigenworms_matrix) + t_m
thetas.append(theta)
thetas = np.array(thetas)
# calculate score and associated skeleton
results_to_score = BaseResults(theta=thetas)
results_scoring(results_to_score, scoring_data_manager)
skeletons = results_to_score.skeletons[:, 0]
scores = results_to_score.scores[:, 0]
# draw skeleton on top of image
color = (0, 255, 0)
image_filename_format = "frame_{{:0{}d}}_score_{{:.2f}}.png".format(
len(str(skeletons.shape[0])), len(str(len(scores))))
with dataset.frames_dataset.open(video_name) as frames:
for index, (frame, skel,
score) in enumerate(zip(frames, skeletons, scores)):
frame_color = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
draw_skeleton(frame_color, skel, color, color)
cv2.imwrite(
os.path.join(out_dir,
image_filename_format.format(index, score)),
frame_color)
# cleanup
shutil.rmtree(temp_dir)
import numpy as np
import matplotlib.pyplot as plt
thetas_onno = np.load("data/AQ2934_theta_onno.npy")
thetas_wp = np.load("data/AQ2934_theta_wp.npy")
from wormpose.pose.eigenworms import load_eigenworms_matrix, theta_to_modes, modes_to_theta
eigenworms_matrix = load_eigenworms_matrix("EigenWorms.csv")
from wormpose.pose.centerline import flip_theta
from wormpose.pose.distance_metrics import angle_distance
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)
def mode_dist(m1, m2):
return np.abs(m1 - m2)
MODE_THRESHOLD = 12
all_mode_errors = []
for index, (theta_wp, theta_onno) in enumerate(zip(thetas_wp, thetas_onno)):
m_wp, t_wp = convert(theta_wp)