File: image_helpers.py

package info (click to toggle)
python-pyface 8.0.0-5
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 13,944 kB
  • sloc: python: 54,107; makefile: 82
file content (174 lines) | stat: -rw-r--r-- 4,532 bytes parent folder | download
1
2
3
4
5
6
7
8
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
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
# (C) Copyright 2005-2023 Enthought, Inc., Austin, TX
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in LICENSE.txt and may be redistributed only under
# the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!

""" Helper functions for working with images

This module provides helper functions for converting between numpy arrays
and Qt wx.Images, as well as between the various image types in a standardized
way.
"""

from enum import IntEnum

import wx


class ScaleMode(IntEnum):
    fast = wx.IMAGE_QUALITY_NORMAL
    smooth = wx.IMAGE_QUALITY_HIGH


class AspectRatio(IntEnum):
    ignore = 0
    keep_constrain = 1
    keep_expand = 2


def image_to_bitmap(image):
    """ Convert a wx.Image to a wx.Bitmap.

    Parameters
    ----------
    image : wx.Image
        The wx.Image to convert.

    Return
    ------
    bitmap : wx.Bitmap
        The corresponding wx.Bitmap.
    """
    return image.ConvertToBitmap()


def bitmap_to_image(bitmap):
    """ Convert a wx.Bitmap to a wx.Image.

    Parameters
    ----------
    bitmap : wx.Bitmap
        The wx.Bitmap to convert.

    Return
    ------
    image : wx.Image
        The corresponding wx.Image.
    """
    return bitmap.ConvertToImage()


def bitmap_to_icon(bitmap):
    """ Convert a wx.Bitmap to a wx.Icon.

    Parameters
    ----------
    bitmap : wx.Bitmap
        The wx.Bitmap to convert.

    Return
    ------
    icon : wx.Icon
        The corresponding wx.Icon.
    """
    return wx.Icon(bitmap)


def resize_image(image, size, aspect_ratio=AspectRatio.ignore,
                 mode=ScaleMode.fast):
    """ Resize a toolkit image to the given size. """
    image_size = image.GetSize()
    width, height = _get_size_for_aspect_ratio(image_size, size, aspect_ratio)
    return image.Scale(width, height, mode)


def resize_bitmap(bitmap, size, aspect_ratio=AspectRatio.ignore,
                  mode=ScaleMode.fast):
    """ Resize a toolkit bitmap to the given size. """
    image = bitmap_to_image(bitmap)
    image = resize_image(image, size, aspect_ratio, mode)
    return image_to_bitmap(image)


def image_to_array(image):
    """ Convert a wx.Image to a numpy array.

    This copies the data returned from wx.

    Parameters
    ----------
    image : wx.Image
        The wx.Image that we want to extract the values from.  The format must
        be either RGB32 or ARGB32.

    Return
    ------
    array : ndarray
        An N x M x 4 array of unsigned 8-bit ints as RGBA values.
    """
    import numpy as np

    width, height = image.GetSize()
    rgb_data = np.array(image.GetData(), dtype='uint8').reshape(width, height, 3)
    if image.HasAlpha():
        alpha = np.array(image.GetAlpha(), dtype='uint8').reshape(width, height)
    else:
        alpha = np.full((width, height), 0xff, dtype='uint8')
    array = np.empty(shape=(width, height, 4), dtype='uint8')
    array[:, :, :3] = rgb_data
    array[:, :, 3] = alpha
    return array


def array_to_image(array):
    """ Convert a numpy array to a wx.Image.

    This copies the data before passing it to wx.

    Parameters
    ----------
    array : ndarray
        An N x M x {3, 4} array of unsigned 8-bit ints.  The image
        format is assumed to be RGB or RGBA, based on the shape.

    Return
    ------
    image : wx.Image
        The wx.Image created from the data.
    """
    if array.ndim != 3:
        raise ValueError("Array must be either RGB or RGBA values.")

    height, width, channels = array.shape

    if channels == 3:
        image = wx.Image(width, height, array.tobytes())
    elif channels == 4:
        image = wx.Image(
            width,
            height,
            array[..., :3].tobytes(),
            array[..., 3].tobytes(),
        )
    else:
        raise ValueError("Array must be either RGB or RGBA values.")

    return image


def _get_size_for_aspect_ratio(image_size, size, aspect_ratio):
    width, height = size
    image_width, image_height = image_size
    if aspect_ratio != AspectRatio.ignore:
        if aspect_ratio == AspectRatio.keep_constrain:
            scale = min(width/image_width, height/image_height)
        else:
            scale = max(width/image_width, height/image_height)
        width = int(round(scale * image_width))
        height = int(round(scale * image_height))
    return (width, height)