Omnibus commited on
Commit
71aa58b
1 Parent(s): 9890778

Create transforms.py

Browse files
Files changed (1) hide show
  1. transforms.py +171 -0
transforms.py ADDED
@@ -0,0 +1,171 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """A library for describing and applying affine transforms to PIL images."""
2
+ import numpy as np
3
+ import PIL.Image
4
+
5
+
6
+ class RGBTransform(object):
7
+ """A description of an affine transformation to an RGB image.
8
+ This class is immutable.
9
+ Methods correspond to matrix left-multiplication/post-application:
10
+ for example,
11
+ RGBTransform().multiply_with(some_color).desaturate()
12
+ describes a transformation where the multiplication takes place first.
13
+ Use rgbt.applied_to(image) to return a converted copy of the given image.
14
+ For example:
15
+ grayish = RGBTransform.desaturate(factor=0.5).applied_to(some_image)
16
+ """
17
+
18
+ def __init__(self, matrix=None):
19
+ self._matrix = matrix if matrix is not None else np.eye(4)
20
+
21
+ def _then(self, operation):
22
+ return RGBTransform(np.dot(_embed44(operation), self._matrix))
23
+
24
+ def desaturate(self, factor=1.0, weights=(0.299, 0.587, 0.114)):
25
+ """Desaturate an image by the given amount.
26
+ A factor of 1.0 will make the image completely gray;
27
+ a factor of 0.0 will leave the image unchanged.
28
+ The weights represent the relative contributions of each channel.
29
+ They should be a 1-by-3 array-like object (tuple, list, np.array).
30
+ In most cases, their values should sum to 1.0
31
+ (otherwise, the transformation will cause the image
32
+ to get lighter or darker).
33
+ """
34
+ weights = _to_rgb(weights, "weights")
35
+
36
+ # tile: [wr, wg, wb] ==> [[wr, wg, wb], [wr, wg, wb], [wr, wg, wb]]
37
+ desaturated_component = factor * np.tile(weights, (3, 1))
38
+ saturated_component = (1 - factor) * np.eye(3)
39
+ operation = desaturated_component + saturated_component
40
+
41
+ return self._then(operation)
42
+
43
+ def multiply_with(self, base_color, factor=1.0):
44
+ """Multiply an image by a constant base color.
45
+ The base color should be a 1-by-3 array-like object
46
+ representing an RGB color in [0, 255]^3 space.
47
+ For example, to multiply with orange,
48
+ the transformation
49
+ RGBTransform().multiply_with((255, 127, 0))
50
+ might be used.
51
+ The factor controls the strength of the multiplication.
52
+ A factor of 1.0 represents straight multiplication;
53
+ other values will be linearly interpolated between
54
+ the identity (0.0) and the straight multiplication (1.0).
55
+ """
56
+ component_vector = _to_rgb(base_color, "base_color") / 255.0
57
+ new_component = factor * np.diag(component_vector)
58
+ old_component = (1 - factor) * np.eye(3)
59
+ operation = new_component + old_component
60
+
61
+ return self._then(operation)
62
+
63
+ def mix_with(self, base_color, factor=1.0):
64
+ """Mix an image by a constant base color.
65
+ The base color should be a 1-by-3 array-like object
66
+ representing an RGB color in [0, 255]^3 space.
67
+ For example, to mix with orange,
68
+ the transformation
69
+ RGBTransform().mix_with((255, 127, 0))
70
+ might be used.
71
+ The factor controls the strength of the color to be added.
72
+ If the factor is 1.0, all pixels will be exactly the new color;
73
+ if it is 0.0, the pixels will be unchanged.
74
+ """
75
+ base_color = _to_rgb(base_color, "base_color")
76
+ operation = _embed44((1 - factor) * np.eye(3))
77
+ operation[:3, 3] = factor * base_color
78
+
79
+ return self._then(operation)
80
+
81
+ def get_matrix(self):
82
+ """Get the underlying 3-by-4 matrix for this affine transform."""
83
+ return self._matrix[:3, :]
84
+
85
+ def applied_to(self, image):
86
+ """Apply this transformation to a copy of the given RGB* image.
87
+ The image should be a PIL image with at least three channels.
88
+ Specifically, the RGB and RGBA modes are both supported, but L is not.
89
+ Any channels past the first three will pass through unchanged.
90
+ The original image will not be modified;
91
+ a new image of the same mode and dimensions will be returned.
92
+ """
93
+
94
+ # PIL.Image.convert wants the matrix as a flattened 12-tuple.
95
+ # (The docs claim that they want a 16-tuple, but this is wrong;
96
+ # cf. _imaging.c:767 in the PIL 1.1.7 source.)
97
+ matrix = tuple(self.get_matrix().flatten())
98
+
99
+ channel_names = image.getbands()
100
+ channel_count = len(channel_names)
101
+ if channel_count < 3:
102
+ raise ValueError("Image must have at least three channels!")
103
+ elif channel_count == 3:
104
+ return image.convert('RGB', matrix)
105
+ else:
106
+ # Probably an RGBA image.
107
+ # Operate on the first three channels (assuming RGB),
108
+ # and tack any others back on at the end.
109
+ channels = list(image.split())
110
+ rgb = PIL.Image.merge('RGB', channels[:3])
111
+ transformed = rgb.convert('RGB', matrix)
112
+ new_channels = transformed.split()
113
+ channels[:3] = new_channels
114
+ return PIL.Image.merge(''.join(channel_names), channels)
115
+
116
+ def applied_to_pixel(self, color):
117
+ """Apply this transformation to a single RGB* pixel.
118
+ In general, you want to apply a transformation to an entire image.
119
+ But in the special case where you know that the image is all one color,
120
+ you can save cycles by just applying the transformation to that color
121
+ and then constructing an image of the desired size.
122
+ For example, in the result of the following code,
123
+ image1 and image2 should be identical:
124
+ rgbt = create_some_rgb_tranform()
125
+ white = (255, 255, 255)
126
+ size = (100, 100)
127
+ image1 = rgbt.applied_to(PIL.Image.new("RGB", size, white))
128
+ image2 = PIL.Image.new("RGB", size, rgbt.applied_to_pixel(white))
129
+ The construction of image2 will be faster for two reasons:
130
+ first, only one PIL image is created; and
131
+ second, the transformation is only applied once.
132
+ The input must have at least three channels;
133
+ the first three channels will be interpreted as RGB,
134
+ and any other channels will pass through unchanged.
135
+ To match the behavior of PIL,
136
+ the values of the resulting pixel will be rounded (not truncated!)
137
+ to the nearest whole number.
138
+ """
139
+ color = tuple(color)
140
+ channel_count = len(color)
141
+ extra_channels = tuple()
142
+ if channel_count < 3:
143
+ raise ValueError("Pixel must have at least three channels!")
144
+ elif channel_count > 3:
145
+ color, extra_channels = color[:3], color[3:]
146
+
147
+ color_vector = np.array(color + (1, )).reshape(4, 1)
148
+ result_vector = np.dot(self._matrix, color_vector)
149
+ result = result_vector.flatten()[:3]
150
+
151
+ full_result = tuple(result) + extra_channels
152
+ rounded = tuple(int(round(x)) for x in full_result)
153
+
154
+ return rounded
155
+
156
+
157
+ def _embed44(matrix):
158
+ """Embed a 4-by-4 or smaller matrix in the upper-left of I_4."""
159
+ result = np.eye(4)
160
+ r, c = matrix.shape
161
+ result[:r, :c] = matrix
162
+ return result
163
+
164
+
165
+ def _to_rgb(thing, name="input"):
166
+ """Convert an array-like object to a 1-by-3 numpy array, or fail."""
167
+ thing = np.array(thing)
168
+ assert thing.shape == (3, ), (
169
+ "Expected %r to be a length-3 array-like object, but found shape %s" %
170
+ (name, thing.shape))
171
+ return thing