Skip to content

Image

ImageSink

Source code in supervision/utils/image.py 
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class ImageSink:
    def __init__(
        self,
        target_dir_path: str,
        overwrite: bool = False,
        image_name_pattern: str = "image_{:05d}.png",
    ):
        """
        Initialize a context manager for saving images.

        Args:
            target_dir_path (str): The target directory where images will be saved.
            overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
            image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".

        Examples:
            ```python
            >>> import supervision as sv

            >>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
            ...     for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
            ...         sink.save_image(image=image)
            ```
        """
        self.target_dir_path = target_dir_path
        self.overwrite = overwrite
        self.image_name_pattern = image_name_pattern
        self.image_count = 0

    def __enter__(self):
        if os.path.exists(self.target_dir_path):
            if self.overwrite:
                shutil.rmtree(self.target_dir_path)
                os.makedirs(self.target_dir_path)
        else:
            os.makedirs(self.target_dir_path)

        return self

    def save_image(self, image: np.ndarray, image_name: Optional[str] = None):
        """
        Save a given image in the target directory.

        Args:
            image (np.ndarray): The image to be saved.
            image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
        """
        if image_name is None:
            image_name = self.image_name_pattern.format(self.image_count)

        image_path = os.path.join(self.target_dir_path, image_name)
        cv2.imwrite(image_path, image)
        self.image_count += 1

    def __exit__(self, exc_type, exc_value, exc_traceback):
        pass

__init__(target_dir_path, overwrite=False, image_name_pattern='image_{:05d}.png')

Initialize a context manager for saving images.

Parameters:

Name Type Description Default
target_dir_path str

The target directory where images will be saved.

 required
overwrite bool

Whether to overwrite the existing directory. Defaults to False.

 False
image_name_pattern str

The image file name pattern. Defaults to "image_{:05d}.png".

 'image_{:05d}.png'

Examples:

>>> import supervision as sv

>>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
...     for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
...         sink.save_image(image=image)
Source code in supervision/utils/image.py 
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
def __init__(
    self,
    target_dir_path: str,
    overwrite: bool = False,
    image_name_pattern: str = "image_{:05d}.png",
):
    """
    Initialize a context manager for saving images.

    Args:
        target_dir_path (str): The target directory where images will be saved.
        overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
        image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".

    Examples:
        ```python
        >>> import supervision as sv

        >>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
        ...     for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
        ...         sink.save_image(image=image)
        ```
    """
    self.target_dir_path = target_dir_path
    self.overwrite = overwrite
    self.image_name_pattern = image_name_pattern
    self.image_count = 0

save_image(image, image_name=None)

Save a given image in the target directory.

Parameters:

Name Type Description Default
image ndarray

The image to be saved.

 required
image_name str

The name to use for the saved image. If not provided, a name will be generated using the image_name_pattern.

 None
Source code in supervision/utils/image.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def save_image(self, image: np.ndarray, image_name: Optional[str] = None):
    """
    Save a given image in the target directory.

    Args:
        image (np.ndarray): The image to be saved.
        image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
    """
    if image_name is None:
        image_name = self.image_name_pattern.format(self.image_count)

    image_path = os.path.join(self.target_dir_path, image_name)
    cv2.imwrite(image_path, image)
    self.image_count += 1

crop

Crops the given image based on the given bounding box.

Parameters:

Name Type Description Default
image ndarray

The image to be cropped, represented as a numpy array.

 required
xyxy ndarray

A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).

 required

Returns:

Type Description
ndarray

The cropped image as a numpy array.

Examples:

>>> import supervision as sv

>>> detection = sv.Detections(...)
>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
...     for xyxy in detection.xyxy:
...         cropped_image = sv.crop(image=image, xyxy=xyxy)
...         sink.save_image(image=image)
Source code in supervision/utils/image.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
    """
    Crops the given image based on the given bounding box.

    Args:
        image (np.ndarray): The image to be cropped, represented as a numpy array.
        xyxy (np.ndarray): A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).

    Returns:
        (np.ndarray): The cropped image as a numpy array.

    Examples:
        ```python
        >>> import supervision as sv

        >>> detection = sv.Detections(...)
        >>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
        ...     for xyxy in detection.xyxy:
        ...         cropped_image = sv.crop(image=image, xyxy=xyxy)
        ...         sink.save_image(image=image)
        ```
    """

    xyxy = np.round(xyxy).astype(int)
    x1, y1, x2, y2 = xyxy
    cropped_img = image[y1:y2, x1:x2]
    return cropped_img