Skip to content

Commit ab2b35c

Browse files
committed
optimise _setSat
1 parent df8de40 commit ab2b35c

File tree

3 files changed

+62
-78
lines changed

3 files changed

+62
-78
lines changed

blendmodes/blend.py

Lines changed: 41 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -193,18 +193,14 @@ def _setLum(originalColours: np.ndarray, newLuminosity: np.ndarray) -> np.ndarra
193193
maxMask = maxColours > 1
194194

195195
# Apply min correction
196-
_colours[minMask] = (
197-
_luminosity[minMask, None] +
198-
((_colours[minMask] - _luminosity[minMask, None]) * _luminosity[minMask, None]) /
199-
(_luminosity[minMask, None] - minColours[minMask, None])
200-
)
196+
_colours[minMask] = _luminosity[minMask, None] + (
197+
(_colours[minMask] - _luminosity[minMask, None]) * _luminosity[minMask, None]
198+
) / (_luminosity[minMask, None] - minColours[minMask, None])
201199

202200
# Apply max correction
203-
_colours[maxMask] = (
204-
_luminosity[maxMask, None] +
205-
((_colours[maxMask] - _luminosity[maxMask, None]) * (1 - _luminosity[maxMask, None])) /
206-
(maxColours[maxMask, None] - _luminosity[maxMask, None])
207-
)
201+
_colours[maxMask] = _luminosity[maxMask, None] + (
202+
(_colours[maxMask] - _luminosity[maxMask, None]) * (1 - _luminosity[maxMask, None])
203+
) / (maxColours[maxMask, None] - _luminosity[maxMask, None])
208204

209205
return _colours
210206

@@ -219,38 +215,42 @@ def _sat(colours: np.ndarray) -> np.ndarray:
219215

220216

221217
def _setSat(originalColours: np.ndarray, newSaturation: np.ndarray) -> np.ndarray:
222-
"""Set a new saturation value for the matrix of color.
223-
224-
The current implementation cannot be vectorized in an efficient manner,
225-
so it is very slow,
226-
O(m*n) at least. This might be able to be improved with openCL if that is
227-
the direction that the lib takes.
228-
:param c: x by x by 3 matrix of rgb color components of pixels
229-
:param s: int of the new saturation value for the matrix
230-
:return: x by x by 3 matrix of luminosity of pixels
231-
"""
218+
"""Set a new saturation value for the matrix of color."""
232219
_colours = originalColours.copy()
233-
for i in range(_colours.shape[0]):
234-
for j in range(_colours.shape[1]):
235-
_colour = _colours[i][j]
236-
minI = 0
237-
midI = 1
238-
maxI = 2
239-
if _colour[midI] < _colour[minI]:
240-
minI, midI = midI, minI
241-
if _colour[maxI] < _colour[midI]:
242-
midI, maxI = maxI, midI
243-
if _colour[midI] < _colour[minI]:
244-
minI, midI = midI, minI
245-
if _colour[maxI] - _colour[minI] > 0.0:
246-
_colours[i][j][midI] = ((_colour[midI] - _colour[minI]) * newSaturation[i, j]) / (
247-
_colour[maxI] - _colour[minI]
248-
)
249-
_colours[i][j][maxI] = newSaturation[i, j]
250-
else:
251-
_colours[i][j][midI] = 0
252-
_colours[i][j][maxI] = 0
253-
_colours[i][j][minI] = 0
220+
221+
# Sort each pixel's color channels to find min, mid, and max
222+
sorted_indices = np.argsort(_colours, axis=2)
223+
minI = sorted_indices[:, :, 0]
224+
midI = sorted_indices[:, :, 1]
225+
maxI = sorted_indices[:, :, 2]
226+
227+
# Extract min, mid, max values
228+
minColours = np.take_along_axis(_colours, minI[..., None], axis=2).squeeze()
229+
midColours = np.take_along_axis(_colours, midI[..., None], axis=2).squeeze()
230+
maxColours = np.take_along_axis(_colours, maxI[..., None], axis=2).squeeze()
231+
232+
# Compute scaling factor
233+
rangeColours = maxColours - minColours
234+
nonzeroMask = rangeColours > 0
235+
236+
# Apply saturation scaling
237+
midColours[nonzeroMask] = (
238+
(midColours[nonzeroMask] - minColours[nonzeroMask]) * newSaturation[nonzeroMask]
239+
) / rangeColours[nonzeroMask]
240+
maxColours[nonzeroMask] = newSaturation[nonzeroMask]
241+
242+
# Zero out mid and max when rangeColours == 0
243+
midColours[~nonzeroMask] = 0
244+
maxColours[~nonzeroMask] = 0
245+
246+
# Set min channel to zero
247+
minColours.fill(0)
248+
249+
# Reassemble the color matrix
250+
np.put_along_axis(_colours, minI[..., None], minColours[..., None], axis=2)
251+
np.put_along_axis(_colours, midI[..., None], midColours[..., None], axis=2)
252+
np.put_along_axis(_colours, maxI[..., None], maxColours[..., None], axis=2)
253+
254254
return _colours
255255

256256

documentation/reference/blendmodes/blend.md

Lines changed: 16 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ def _lum(colours: np.ndarray) -> np.ndarray: ...
7272

7373
## _sat
7474

75-
[Show source in blend.py:205](../../../blendmodes/blend.py#L205)
75+
[Show source in blend.py:208](../../../blendmodes/blend.py#L208)
7676

7777
Saturation.
7878

@@ -108,24 +108,10 @@ def _setLum(originalColours: np.ndarray, newLuminosity: np.ndarray) -> np.ndarra
108108

109109
## _setSat
110110

111-
[Show source in blend.py:214](../../../blendmodes/blend.py#L214)
111+
[Show source in blend.py:217](../../../blendmodes/blend.py#L217)
112112

113113
Set a new saturation value for the matrix of color.
114114

115-
The current implementation cannot be vectorized in an efficient manner,
116-
so it is very slow,
117-
O(m*n) at least. This might be able to be improved with openCL if that is
118-
the direction that the lib takes.
119-
120-
#### Arguments
121-
122-
- `c` - x by x by 3 matrix of rgb color components of pixels
123-
- `s` - int of the new saturation value for the matrix
124-
125-
#### Returns
126-
127-
x by x by 3 matrix of luminosity of pixels
128-
129115
#### Signature
130116

131117
```python
@@ -150,7 +136,7 @@ def additive(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
150136

151137
## alpha_comp_shell
152138

153-
[Show source in blend.py:631](../../../blendmodes/blend.py#L631)
139+
[Show source in blend.py:638](../../../blendmodes/blend.py#L638)
154140

155141
Implement common transformations occurring in any blend or composite mode.
156142

@@ -170,7 +156,7 @@ def alpha_comp_shell(
170156

171157
## blend
172158

173-
[Show source in blend.py:393](../../../blendmodes/blend.py#L393)
159+
[Show source in blend.py:400](../../../blendmodes/blend.py#L400)
174160

175161
Blend pixels.
176162

@@ -223,7 +209,7 @@ def blend(
223209

224210
## blendLayers
225211

226-
[Show source in blend.py:460](../../../blendmodes/blend.py#L460)
212+
[Show source in blend.py:467](../../../blendmodes/blend.py#L467)
227213

228214
Blend two layers (background, and foreground).
229215

@@ -275,7 +261,7 @@ def blendLayers(
275261

276262
## blendLayersArray
277263

278-
[Show source in blend.py:510](../../../blendmodes/blend.py#L510)
264+
[Show source in blend.py:517](../../../blendmodes/blend.py#L517)
279265

280266
Blend two layers (background, and foreground).
281267

@@ -331,7 +317,7 @@ def blendLayersArray(
331317

332318
## colour
333319

334-
[Show source in blend.py:260](../../../blendmodes/blend.py#L260)
320+
[Show source in blend.py:267](../../../blendmodes/blend.py#L267)
335321

336322
BlendType.COLOUR.
337323

@@ -387,7 +373,7 @@ def darken(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
387373

388374
## destatop
389375

390-
[Show source in blend.py:321](../../../blendmodes/blend.py#L321)
376+
[Show source in blend.py:328](../../../blendmodes/blend.py#L328)
391377

392378
Place the layer below above the 'layer above' in places where the 'layer above' exists...
393379

@@ -408,7 +394,7 @@ def destatop(
408394

409395
## destin
410396

411-
[Show source in blend.py:270](../../../blendmodes/blend.py#L270)
397+
[Show source in blend.py:277](../../../blendmodes/blend.py#L277)
412398

413399
'clip' composite mode.
414400

@@ -437,7 +423,7 @@ def destin(
437423

438424
## destout
439425

440-
[Show source in blend.py:298](../../../blendmodes/blend.py#L298)
426+
[Show source in blend.py:305](../../../blendmodes/blend.py#L305)
441427

442428
Reverse 'Clip' composite mode.
443429

@@ -558,7 +544,7 @@ def hardlight(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
558544

559545
## hue
560546

561-
[Show source in blend.py:250](../../../blendmodes/blend.py#L250)
547+
[Show source in blend.py:257](../../../blendmodes/blend.py#L257)
562548

563549
BlendType.HUE.
564550

@@ -572,7 +558,7 @@ def hue(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
572558

573559
## imageFloatToInt
574560

575-
[Show source in blend.py:375](../../../blendmodes/blend.py#L375)
561+
[Show source in blend.py:382](../../../blendmodes/blend.py#L382)
576562

577563
Convert a numpy array representing an image to an array of ints.
578564

@@ -596,7 +582,7 @@ def imageFloatToInt(image: np.ndarray) -> np.ndarray: ...
596582

597583
## imageIntToFloat
598584

599-
[Show source in blend.py:360](../../../blendmodes/blend.py#L360)
585+
[Show source in blend.py:367](../../../blendmodes/blend.py#L367)
600586

601587
Convert a numpy array representing an image to an array of floats.
602588

@@ -634,7 +620,7 @@ def lighten(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
634620

635621
## luminosity
636622

637-
[Show source in blend.py:265](../../../blendmodes/blend.py#L265)
623+
[Show source in blend.py:272](../../../blendmodes/blend.py#L272)
638624

639625
BlendType.LUMINOSITY.
640626

@@ -732,7 +718,7 @@ def reflect(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
732718

733719
## saturation
734720

735-
[Show source in blend.py:255](../../../blendmodes/blend.py#L255)
721+
[Show source in blend.py:262](../../../blendmodes/blend.py#L262)
736722

737723
BlendType.SATURATION.
738724

@@ -774,7 +760,7 @@ def softlight(background: np.ndarray, foreground: np.ndarray) -> np.ndarray: ...
774760

775761
## srcatop
776762

777-
[Show source in blend.py:342](../../../blendmodes/blend.py#L342)
763+
[Show source in blend.py:349](../../../blendmodes/blend.py#L349)
778764

779765
Place the layer below above the 'layer above' in places where the 'layer above' exists.
780766

tests/test_perf.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,12 @@
1111
THISDIR = Path(__file__).resolve().parent
1212

1313

14-
15-
1614
@pytest.mark.parametrize("blend_mode", random.choices(list(BlendType), k=200))
1715
def test_blend_modes(blend_mode: BlendType) -> None:
18-
"""Test blend modes with random selection."""
19-
background = Image.open(THISDIR / "data" / "background.png")
20-
foreground = Image.open(THISDIR / "data" / "foreground.png")
16+
"""Test blend modes with random selection."""
17+
background = Image.open(THISDIR / "data" / "background.png")
18+
foreground = Image.open(THISDIR / "data" / "foreground.png")
2119

22-
result = blendLayers(background, foreground, blend_mode)
20+
result = blendLayers(background, foreground, blend_mode)
2321

24-
assert result is not None # Ensure we get a valid image back
22+
assert result is not None # Ensure we get a valid image back

0 commit comments

Comments
 (0)