def load_dataset(
dataset_loader: str,
dataset_path: str,
selected_video_names: Optional[List[str]] = None,
resize_options: ResizeOptions = None,
) -> Dataset:
for entry_point in pkg_resources.iter_entry_points("worm_dataset_loaders"):
if entry_point.name == dataset_loader:
module = entry_point.load()
frames_dataset_class = module.FramesDataset
features_dataset_class = module.FeaturesDataset
frame_preprocessing_class = module.FramePreprocessing
frames_dataset, features_dataset, video_names = _load_dataset(
frames_dataset_class,
features_dataset_class,
dataset_path,
selected_video_names,
)
frame_preprocessing = frame_preprocessing_class()
image_shape = calculate_crop_window_size(features_dataset)
if resize_options is not None:
resize_options.update_resize_factor(features_dataset)
if resize_options.resize_factor != 1.0:
# reload frames and features dataset after resizing, also get new image_shape
frames_dataset_class = frames_dataset_resizer(
frames_dataset_class,
resize_factor=resize_options.resize_factor)
features_dataset_class = features_dataset_resizer(
features_dataset_class,
resize_factor=resize_options.resize_factor,
)
frames_dataset, features_dataset, video_names = _load_dataset(
frames_dataset_class,
features_dataset_class,
dataset_path,
selected_video_names,
)
image_shape = resize_options.get_image_shape(
features_dataset)
results_exporter = (module.ResultsExporter(dataset_path)
if hasattr(module, "ResultsExporter") else
_DummyResultsExporter())
return Dataset(
video_names=video_names,
features_dataset=features_dataset,
frames_dataset=frames_dataset,
frame_preprocessing=frame_preprocessing,
image_shape=image_shape,
results_exporter=results_exporter,
)
raise NotImplementedError(
f"Dataset loader: '{dataset_loader}' not found in the package entry points."
)
def __init__(
self,
dataset_loader: str,
dataset_path: str,
postures_generator: Optional[Generator] = None,
video_name: str = None,
**kwargs
):
resize_options = ResizeOptions(**kwargs)
dataset = load_dataset(dataset_loader, dataset_path, resize_options=resize_options, **kwargs)
if postures_generator is None:
postures_generator = PosturesModel().generate()
if video_name is None:
video_name = dataset.video_names[0]
features = dataset.features_dataset[video_name]
self.skeletons = features.skeletons
self.measurements = features.measurements
self.output_image_shape = dataset.image_shape
self.synthetic_dataset = SyntheticDataset(
frame_preprocessing=dataset.frame_preprocessing,
output_image_shape=self.output_image_shape,
enable_random_augmentations=False,
)
skel_is_not_nan = ~np.any(np.isnan(self.skeletons), axis=(1, 2))
self.labelled_indexes = np.where(skel_is_not_nan)[0]
if len(self.labelled_indexes) == 0:
raise ValueError("No template frames found in the dataset, can't generate synthetic images.")
self.frames_dataset = dataset.frames_dataset
self.video_name = video_name
self.postures_generator = postures_generator
def __init__(self, dataset_loader: str, dataset_path: str, video_name=None, **kwargs):
resize_options = ResizeOptions(**kwargs)
dataset = load_dataset(dataset_loader, dataset_path, resize_options=resize_options)
self.video_name = video_name if video_name is not None else dataset.video_names[0]
self.real_dataset = RealDataset(
frame_preprocessing=dataset.frame_preprocessing,
output_image_shape=dataset.image_shape,
)
self.frames_dataset = dataset.frames_dataset
def _parse_arguments(kwargs: dict):
if kwargs.get("num_samples") is None:
kwargs["num_samples"] = 500
if kwargs.get("work_dir") is None:
kwargs["work_dir"] = default_paths.WORK_DIR
if kwargs.get("theta_dims") is None:
kwargs["theta_dims"] = 100
if kwargs.get("video_names") is None:
kwargs["video_names"] = None
if kwargs.get("save_images") is None:
kwargs["save_images"] = False
if kwargs.get("random_seed") is None:
kwargs["random_seed"] = None
kwargs["resize_options"] = ResizeOptions(**kwargs)
_log_parameters(logger.info, kwargs)
return Namespace(**kwargs)
def _parse_arguments(kwargs: dict):
if kwargs.get("num_process") is None:
kwargs["num_process"] = os.cpu_count()
if kwargs.get("temp_dir") is None:
kwargs["temp_dir"] = tempfile.gettempdir()
if kwargs.get("num_train_samples") is None:
kwargs["num_train_samples"] = int(5e5)
if kwargs.get("num_eval_samples") is None:
kwargs["num_eval_samples"] = int(1e4)
if kwargs.get("work_dir") is None:
kwargs["work_dir"] = default_paths.WORK_DIR
if kwargs.get("postures_generation") is None:
kwargs["postures_generation"] = PosturesModel().generate
if kwargs.get("video_names") is None:
kwargs["video_names"] = None
if kwargs.get("random_seed") is None:
kwargs["random_seed"] = None
kwargs["temp_dir"] = tempfile.mkdtemp(dir=kwargs["temp_dir"])
kwargs["resize_options"] = ResizeOptions(**kwargs)
_log_parameters(logger.info, kwargs)
return Namespace(**kwargs)
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 predict(dataset_path: str, **kwargs):
"""
Use a trained model to predict the centerlines of worm for videos in a dataset
: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)
if args.random_seed is not None:
os.environ["TF_DETERMINISTIC_OPS"] = "1"
random.seed(args.random_seed)
np.random.seed(args.random_seed)
tf.random.set_seed(args.random_seed)
results_root_dir = os.path.join(args.experiment_dir,
default_paths.RESULTS_DIR)
os.makedirs(results_root_dir, exist_ok=True)
config = load_config(args.config)
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),
)
keras_model = tf.keras.models.load_model(args.model_path, compile=False)
results_saver = ResultsSaver(temp_dir=args.temp_dir,
results_root_dir=results_root_dir,
results_filename=RESULTS_FILENAME)
tf_dataset_maker = _make_tf_dataset(
data_generator=PredictDataGenerator(
dataset=dataset,
num_process=args.num_process,
temp_dir=args.temp_dir,
image_shape=config.image_shape,
batch_size=args.batch_size,
),
batch_size=args.batch_size,
image_shape=config.image_shape,
)
results_scoring = ResultsScoring(
frame_preprocessing=dataset.frame_preprocessing,
num_process=args.num_process,
temp_dir=args.temp_dir,
image_shape=config.image_shape,
)
predictor = _Predictor(results_scoring=results_scoring,
keras_model=keras_model)
for video_name in dataset.video_names:
logger.info(f'Processing video: "{video_name}"')
features = dataset.features_dataset[video_name]
template_indexes = features.labelled_indexes
if len(template_indexes) == 0:
logger.error(
f"Can't calculate image metric, there is no labelled frame in the video to use as a template, "
f"stopping analysis for {video_name}.")
continue
original_results, shuffled_results = predictor(
input_frames=tf_dataset_maker(video_name),
num_frames=dataset.num_frames(video_name),
features=features,
scoring_data_manager=ScoringDataManager(
video_name=video_name,
frames_dataset=dataset.frames_dataset,
features=features,
),
)
results = {"original": original_results, "unaligned": shuffled_results}
if _can_resolve_results(
shuffled_results,
video_name=video_name,
score_threshold=args.score_threshold,
):
final_results = resolve_head_tail(
shuffled_results=shuffled_results,
original_results=original_results,
frame_rate=features.frame_rate,
score_threshold=args.score_threshold,
)
results["resolved"] = final_results
_apply_resize_factor(results["resolved"], config.resize_factor)
_apply_resize_factor(results["unaligned"], config.resize_factor)
results_saver.save(results=results, video_name=video_name)
# cleanup
shutil.rmtree(args.temp_dir)
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)