コード例 #1
0
    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
コード例 #2
0
def _write_to_file(
    out_filename: str,
    num_samples: int,
    template_filename: str,
    postures_generation_fn: Callable[[], Generator],
    progress_counter,
    writer: Type[GenericFileWriter],
    synthetic_dataset_args,
    random_seed: Optional[int],
):
    from wormpose.images.synthetic import SyntheticDataset

    seed = int.from_bytes(os.urandom(4), byteorder="little") if random_seed is None else random_seed
    np.random.seed(seed)

    with open(template_filename, "rb") as templates_f:
        templates = pickle.load(templates_f)
    os.remove(template_filename)

    worm_measurements = {}
    for video_name in np.unique(templates.video_names):
        indexes = np.where(templates.video_names == video_name)[0]
        measurements = templates.measurements[indexes]
        average_measurements = np.empty((measurements.shape[1:]), dtype=measurements.dtype)
        for name in measurements.dtype.names:
            average_measurements[0][name] = np.nanmean(measurements[name])
        worm_measurements[video_name] = average_measurements

    synthetic_dataset = SyntheticDataset(**synthetic_dataset_args)

    # preallocate all the template indexes to create the synthetic images (random)
    template_indexes = np.random.randint(len(templates.frames), size=num_samples)

    # preallocate all the choices for the position of the head (random)
    headtail_choice = np.random.choice([0, 1], size=num_samples)

    # preallocate image buffer to draw the synthetic image on
    image_data = np.empty(synthetic_dataset.output_image_shape, dtype=np.uint8)

    postures_gen = postures_generation_fn()
    with writer(out_filename) as synth_data_writer:

        for index, (cur_template_index, cur_headtail_choice) in enumerate(zip(template_indexes, headtail_choice)):
            theta = next(postures_gen)
            label_data = [theta, flip_theta(theta)]
            video_name = templates.video_names[cur_template_index]
            template_frame = templates.frames[cur_template_index]
            template_skeleton = templates.skeletons[cur_template_index]
            template_measurements = worm_measurements[video_name]
            synthetic_dataset.generate(
                theta=label_data[cur_headtail_choice],
                template_frame=template_frame,
                template_skeleton=template_skeleton,
                template_measurements=template_measurements,
                out_image=image_data,
            )
            synth_data_writer.write(locals())
            progress_counter.value = index + 1
コード例 #3
0
class SyntheticSimpleVisualizer(object):
    """
    Utility class to visualize the synthetic images
    """
    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)

        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 generate(self):
        out_image = np.empty(self.output_image_shape, dtype=np.uint8)

        with self.frames_dataset.open(self.video_name) as frames:
            while True:
                theta = next(self.postures_generator)
                random_label_index = np.random.choice(self.labelled_indexes)
                self.synthetic_dataset.generate(
                    theta=theta,
                    template_skeleton=self.skeletons[random_label_index],
                    template_frame=frames[random_label_index],
                    out_image=out_image,
                    template_measurements=self.measurements,
                )
                yield out_image, theta
コード例 #4
0
    def __init__(self, frame_preprocessing: BaseFramePreprocessing,
                 image_shape):

        self.real_dataset = RealDataset(
            frame_preprocessing=frame_preprocessing,
            output_image_shape=image_shape)

        self.synthetic_dataset = SyntheticDataset(
            frame_preprocessing=frame_preprocessing,
            output_image_shape=image_shape,
            enable_random_augmentations=False,
        )
        self.last_synth_image = np.empty(image_shape, np.uint8)
        self.last_real_image = None
コード例 #5
0
class CenterlineAccuracyCheck(object):
    """
    Class that performs comparison (image similarity) to assess how a centerline accurately represents a real image.

    It first preprocesses the original real image with a BaseFramePreprocessing class
    (crop and set the background pixels to a uniform color).
    Then, it creates a synthetic image representing the centerline theta, using a provided template image.
    Typically, the template image was chosen to be the closest labelled image in time to the real image.
    The synthetic image is cropped to fit the worm, in order to apply a template matching function between
    the real image (full size) and the synthetic image (smaller)
    The result is an image similarity value and the synthetic image skeleton coordinates.
    """
    def __init__(self, frame_preprocessing: BaseFramePreprocessing,
                 image_shape):

        self.real_dataset = RealDataset(
            frame_preprocessing=frame_preprocessing,
            output_image_shape=image_shape)

        self.synthetic_dataset = SyntheticDataset(
            frame_preprocessing=frame_preprocessing,
            output_image_shape=image_shape,
            enable_random_augmentations=False,
        )
        self.last_synth_image = np.empty(image_shape, np.uint8)
        self.last_real_image = None

    def __call__(
        self,
        theta,
        template_skeleton,
        template_frame,
        template_measurements,
        real_frame_orig,
    ):
        if np.any(np.isnan(theta)):
            score = np.nan
            synth_skel = np.full_like(template_skeleton, np.nan)
            return score, synth_skel

        self.last_real_image, skel_offset = self.real_dataset.process_frame(
            real_frame_orig)
        cur_bg_color, synth_skel = self.synthetic_dataset.generate(
            theta,
            template_frame=template_frame,
            template_skeleton=template_skeleton,
            out_image=self.last_synth_image,
            template_measurements=template_measurements,
        )

        # Crop the synthetic image to the object of interest before doing the image comparison,
        # we don't need the full image with all the background, still keep a little padding around the worm.
        bounding_box = fit_bounding_box_to_worm(self.last_synth_image,
                                                cur_bg_color)
        np.subtract(synth_skel, (bounding_box[1].start, bounding_box[0].start),
                    out=synth_skel)
        cropped_synth_image = self.last_synth_image[bounding_box]

        # Perform the image comparison between the real image and the reconstructed synthetic image cropped.
        # This gives a heat map, we are interested in the maximum value of the heatmap and its location.
        # This maximum score gives an estimation of the confidence of the prediction.
        score, score_loc = image_scoring.calculate_similarity(
            source_image=self.last_real_image,
            template_image=cropped_synth_image)

        # Using the heatmap maximum coordinates, we can transform the coordinates of the reconstructed skeleton
        # to the original image coordinates system.
        synth_skel += score_loc
        synth_skel += np.array([skel_offset[0], skel_offset[1]])

        return score, synth_skel