Skip to content

image_viewer

ImageViewer

Bases: Trace

A trace that interrupts your training in order to display images on the screen.

This class is useful primarily for Jupyter Notebook, or for debugging purposes.

Parameters:

Name Type Description Default
inputs Union[str, Sequence[str]]

Key(s) of images to be displayed.

 required
mode Union[None, str, Iterable[str]]

What mode(s) to execute this Trace in. For example, "train", "eval", "test", or "infer". To execute regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument like "!infer" or "!train".

 ('eval', 'test')
interactive bool

Whether the image should be interactive. This is False by default to reduce jupyter file size.

 False
Source code in fastestimator/fastestimator/trace/io/image_viewer.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@traceable()
class ImageViewer(Trace):
    """A trace that interrupts your training in order to display images on the screen.

    This class is useful primarily for Jupyter Notebook, or for debugging purposes.

    Args:
        inputs: Key(s) of images to be displayed.
        mode: What mode(s) to execute this Trace in. For example, "train", "eval", "test", or "infer". To execute
            regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
            like "!infer" or "!train".
        interactive: Whether the image should be interactive. This is False by default to reduce jupyter file size.
    """
    def __init__(self,
                 inputs: Union[str, Sequence[str]],
                 mode: Union[None, str, Iterable[str]] = ("eval", "test"),
                 interactive: bool = False,
                 ) -> None:
        super().__init__(inputs=inputs, mode=mode)
        self.interactive = interactive

    def on_epoch_end(self, data: Data) -> None:
        self._display_images(data)

    def on_end(self, data: Data) -> None:
        self._display_images(data)

    def _display_images(self, data: Data) -> None:
        """A method to render images to the screen.

        Args:
            data: Data possibly containing images to render.
        """
        for key in self.inputs:
            if key in data:
                imgs = data[key]
                if isinstance(imgs, Display):
                    imgs.show(interactive=self.interactive)
                elif isinstance(imgs, Summary):
                    visualize_logs([imgs])
                elif isinstance(imgs, (list, tuple)) and all([isinstance(img, Summary) for img in imgs]):
                    visualize_logs(imgs)
                else:
                    for idx, img in enumerate(imgs):
                        fig = ImageDisplay(image=img, title="{}_{}".format(key, idx))
                        fig.show(interactive=self.interactive)