segment.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. #!/usr/bin/env python3
  2. """
  3. WLED Segment class for Raspberry Pi
  4. Manages LED strip segments and their effects
  5. """
  6. import time
  7. from typing import List, Callable, Optional
  8. from .utils.colors import *
  9. from .utils.palettes import color_from_palette, get_palette
  10. class Segment:
  11. """LED segment with effect support"""
  12. def __init__(self, pixels, start: int, stop: int, is_rgbw: bool = False):
  13. """
  14. Initialize a segment
  15. pixels: neopixel.NeoPixel object
  16. start: starting LED index
  17. stop: ending LED index (exclusive)
  18. is_rgbw: True for RGBW strips (SK6812), writes 4-tuples instead of 3
  19. """
  20. self.pixels = pixels
  21. self.start = start
  22. self.stop = stop
  23. self.is_rgbw = is_rgbw
  24. self.length = stop - start
  25. # Colors (up to 3 colors like WLED)
  26. self.colors = [0x00FF0000, 0x000000FF, 0x0000FF00] # Red, Blue, Green defaults
  27. # Effect parameters
  28. self.speed = 128 # 0-255
  29. self.intensity = 128 # 0-255
  30. self.palette_id = 0 # Palette ID
  31. self.custom1 = 0 # Custom parameter 1
  32. self.custom2 = 0 # Custom parameter 2
  33. self.custom3 = 0 # Custom parameter 3
  34. # Runtime state (like SEGENV in WLED)
  35. self.call = 0 # Number of times effect has been called
  36. self.step = 0 # Effect step counter
  37. self.aux0 = 0 # Auxiliary variable 0
  38. self.aux1 = 0 # Auxiliary variable 1
  39. self.next_time = 0 # Next time to run effect (ms)
  40. self.data = [] # Effect data storage
  41. # Timing
  42. self._start_time = time.time() * 1000 # Convert to milliseconds
  43. def get_pixel_color(self, i: int) -> int:
  44. """Get color of pixel at index i (segment-relative)"""
  45. if i < 0 or i >= self.length:
  46. return 0
  47. actual_idx = self.start + i
  48. color = self.pixels[actual_idx]
  49. # Convert from neopixel (R,G,B) or (G,R,B) to 32-bit
  50. if isinstance(color, tuple):
  51. if len(color) == 3:
  52. return (color[0] << 16) | (color[1] << 8) | color[2]
  53. elif len(color) == 4:
  54. return (color[3] << 24) | (color[0] << 16) | (color[1] << 8) | color[2]
  55. return 0
  56. def set_pixel_color(self, i: int, color: int):
  57. """Set color of pixel at index i (segment-relative)"""
  58. if i < 0 or i >= self.length:
  59. return
  60. actual_idx = self.start + i
  61. r = get_r(color)
  62. g = get_g(color)
  63. b = get_b(color)
  64. if self.is_rgbw:
  65. w = get_w(color)
  66. self.pixels[actual_idx] = (r, g, b, w)
  67. else:
  68. self.pixels[actual_idx] = (r, g, b)
  69. def fill(self, color: int):
  70. """Fill entire segment with color"""
  71. for i in range(self.length):
  72. self.set_pixel_color(i, color)
  73. def fade_out(self, amount: int):
  74. """Fade all pixels in segment toward black"""
  75. for i in range(self.length):
  76. current = self.get_pixel_color(i)
  77. faded = color_fade(current, amount)
  78. self.set_pixel_color(i, faded)
  79. def blur(self, amount: int):
  80. """Blur/blend pixels with neighbors"""
  81. if amount == 0 or self.length < 3:
  82. return
  83. # Simple blur: blend each pixel with neighbors
  84. temp = [self.get_pixel_color(i) for i in range(self.length)]
  85. for i in range(self.length):
  86. if i == 0:
  87. # First pixel: blend with next
  88. blended = color_blend(temp[i], temp[i + 1], amount)
  89. elif i == self.length - 1:
  90. # Last pixel: blend with previous
  91. blended = color_blend(temp[i], temp[i - 1], amount)
  92. else:
  93. # Middle pixels: blend with both neighbors
  94. left = color_blend(temp[i], temp[i - 1], amount // 2)
  95. blended = color_blend(left, temp[i + 1], amount // 2)
  96. self.set_pixel_color(i, blended)
  97. def now(self) -> int:
  98. """Get current time in milliseconds since segment start"""
  99. return int((time.time() * 1000) - self._start_time)
  100. def color_from_palette(self, index: int, use_index: bool = True,
  101. brightness: int = 255) -> int:
  102. """
  103. Get color from current palette
  104. index: LED index or palette position
  105. use_index: if True, scale index across palette
  106. brightness: 0-255
  107. """
  108. palette = get_palette(self.palette_id)
  109. if use_index and self.length > 1:
  110. # Map LED position to palette position
  111. palette_pos = (index * 255) // (self.length - 1)
  112. else:
  113. palette_pos = index & 0xFF
  114. return color_from_palette(palette, palette_pos, brightness)
  115. def get_color(self, index: int) -> int:
  116. """Get segment color by index (0-2)"""
  117. if 0 <= index < len(self.colors):
  118. return self.colors[index]
  119. return 0
  120. def reset(self):
  121. """Reset segment state"""
  122. self.call = 0
  123. self.step = 0
  124. self.aux0 = 0
  125. self.aux1 = 0
  126. self.next_time = 0
  127. self.data = []