Merge pull request #1 from relvacode/image-downsample

Faster palette extraction through image downsampling

Author: relvacode

README.md

@@ -48,4 +48,14 @@ media_player:
   - platform: lightcast
     name: Office LightCast
     target: area.office
+```
+
+It can take a while to extract the palette of very large images. Use `downsample` to downsample images to a width of `256`. The default is true.
+
+```yaml
+media_player:
+  - platform: lightcast
+    name: Living Room LightCast
+    target: area.living_room
+    # downsample: false
 ```

custom_components/lightcast/color_extract.py

@@ -0,0 +1,43 @@
+import io
+import typing
+import itertools
+
+import colorgram
+from PIL import Image
+
+RGB = typing.Tuple[float, float, float]
+
+MAX_SAMPLE_SIZE = 256
+
+
+def extract_palette(f: io.BytesIO | Image.Image, colors: int, downsample: bool = True) -> list[RGB]:
+    image = f if isinstance(f, Image.Image) else Image.open(f)
+    if image.mode not in ('RGB', 'RGBA', 'RGBa'):
+        image = image.convert('RGB')
+
+    # Downsample image
+    if downsample and (image.width > MAX_SAMPLE_SIZE or image.height > MAX_SAMPLE_SIZE):
+        aspect = MAX_SAMPLE_SIZE / float(image.width)
+        height = int(float(image.height) * aspect)
+        image = image.resize((MAX_SAMPLE_SIZE, height), resample=Image.Resampling.NEAREST)
+
+    extracted_colors = colorgram.extract(image, colors)
+
+    palette = []
+    for i, color in zip(range(colors), itertools.cycle(extracted_colors)):
+        palette.append((color.rgb.r, color.rgb.g, color.rgb.b))
+
+    return palette
+
+
+if __name__ == "__main__":
+    def main():
+        import sys
+
+        with open(sys.argv[1], 'rb') as f:
+            palette = extract_palette(f, int(sys.argv[2]))
+            for color in palette:
+                print(color)
+
+
+    main()

custom_components/lightcast/const.py

@@ -1,3 +1,4 @@
 CONF_TARGET = 'target'
 CONF_NAME = 'name'
 CONF_FILTER_ON = 'filter_on'
+CONF_DOWNSAMPLE = 'downsample'

custom_components/lightcast/media_player.py

@@ -1,8 +1,6 @@
 import logging
 import io
-import itertools
 import aiohttp
-import colorgram
 
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@@ -19,6 +17,7 @@
 
 from . import const
 from .entity_resolver import expand_entities
+from .color_extract import extract_palette
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -35,7 +34,8 @@ async def async_setup_platform(
         hass,
         config[const.CONF_NAME],
         config[const.CONF_TARGET],
-        config.get(const.CONF_FILTER_ON) or False
+        config.get(const.CONF_FILTER_ON) or False,
+        True if (downsample := config.get(const.CONF_DOWNSAMPLE)) is None else downsample,
     )
 
     add_entities([cast_device])
@@ -46,11 +46,12 @@ class LightCastPlayer(MediaPlayerEntity):
     _attr_supported_features = MediaPlayerEntityFeature.PLAY_MEDIA | MediaPlayerEntityFeature.BROWSE_MEDIA
     _attr_device_class = MediaPlayerDeviceClass.TV
 
-    def __init__(self, hass: HomeAssistant, name: str, device_target: str, filter_on: bool) -> None:
+    def __init__(self, hass: HomeAssistant, name: str, device_target: str, filter_on: bool, downsample: bool) -> None:
         self.hass = hass
         self._attr_name = name
         self.device_target = device_target
         self.filter_on = filter_on
+        self.downsample = downsample
 
     async def async_browse_media(self, media_content_type: str | None = None,
                                  media_content_id: str | None = None) -> BrowseMedia:
@@ -95,21 +96,16 @@ async def process_image(self, media_type: str, media_id: str) -> None:
                                    ) as response:
                 response_data = await response.read()
 
-        n_colors = len(valid_entities)
-        _LOGGER.info('Downloaded image, extracting palette of %d colors', n_colors)
+        palette = extract_palette(io.BytesIO(response_data), len(valid_entities), downsample=self.downsample)
 
-        extracted_colors = colorgram.extract(io.BytesIO(response_data), n_colors)
-        _LOGGER.info('Extracted palette of %d colors from image', len(extracted_colors))
-
-        for e, color in zip(valid_entities, itertools.cycle(extracted_colors)):
-            rgb_color = [color.rgb.r, color.rgb.g, color.rgb.b]
-            _LOGGER.info('Set light %s to %s', e.entity_id, rgb_color)
+        for e, color in zip(valid_entities, palette):
+            _LOGGER.info('Set light %s to %s', e.entity_id, color)
             await self.hass.services.async_call(
                 light.DOMAIN,
                 light.SERVICE_TURN_ON,
                 {
                     ATTR_ENTITY_ID: e.entity_id,
-                    light.ATTR_RGB_COLOR: rgb_color
+                    light.ATTR_RGB_COLOR: color
                 }
             )