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)