|
| 1 | +import threading |
| 2 | +import numpy |
| 3 | +import spidev |
| 4 | + |
| 5 | +class SPI_WS2812_LEDStrip: |
| 6 | + def __init__(self, count=8, brightness=255, sequence='GRB', bus=0, device=0): |
| 7 | + self.set_led_type(sequence) |
| 8 | + self.set_led_count(count) |
| 9 | + self.set_led_brightness(brightness) |
| 10 | + self.led_begin(bus, device) |
| 11 | + self.set_all_led_color(0, 0, 0) |
| 12 | + self.lock = threading.Lock() |
| 13 | + |
| 14 | + def set_led_type(self, rgb_type): |
| 15 | + led_types = ['RGB', 'RBG', 'GRB', 'GBR', 'BRG', 'BGR'] |
| 16 | + offsets = [0x06, 0x09, 0x12, 0x21, 0x18, 0x24] |
| 17 | + try: |
| 18 | + idx = led_types.index(rgb_type) |
| 19 | + offset = offsets[idx] |
| 20 | + self.led_red_offset = (offset >> 4) & 3 |
| 21 | + self.led_green_offset = (offset >> 2) & 3 |
| 22 | + self.led_blue_offset = offset & 3 |
| 23 | + except ValueError: |
| 24 | + self.led_red_offset, self.led_green_offset, self.led_blue_offset = 1, 0, 2 |
| 25 | + |
| 26 | + def set_led_count(self, count): |
| 27 | + self.led_count = count |
| 28 | + self.led_color = [0] * (count * 3) |
| 29 | + self.led_original_color = [0] * (count * 3) |
| 30 | + |
| 31 | + def set_led_brightness(self, brightness): |
| 32 | + self.led_brightness = brightness |
| 33 | + for i in range(self.led_count): |
| 34 | + self.set_led_rgb_data(i, [0, 0, 0]) |
| 35 | + |
| 36 | + def set_ledpixel(self, index, r, g, b): |
| 37 | + p = [0, 0, 0] |
| 38 | + p[self.led_red_offset] = round(r * self.led_brightness / 255) |
| 39 | + p[self.led_green_offset] = round(g * self.led_brightness / 255) |
| 40 | + p[self.led_blue_offset] = round(b * self.led_brightness / 255) |
| 41 | + for i, color in enumerate((r, g, b)): |
| 42 | + self.led_original_color[index * 3 + i] = color |
| 43 | + for i in range(3): |
| 44 | + self.led_color[index * 3 + i] = p[i] |
| 45 | + |
| 46 | + def set_led_rgb_data(self, index, color): |
| 47 | + self.set_ledpixel(index, *color) |
| 48 | + |
| 49 | + def set_all_led_color(self, r, g, b): |
| 50 | + for i in range(self.led_count): |
| 51 | + self.set_ledpixel(i, r, g, b) |
| 52 | + self.show() |
| 53 | + |
| 54 | + def led_begin(self, bus=0, device=0): |
| 55 | + self.bus, self.device = bus, device |
| 56 | + try: |
| 57 | + self.spi = spidev.SpiDev() |
| 58 | + self.spi.open(bus, device) |
| 59 | + self.spi.mode = 0 |
| 60 | + self.led_init_state = 1 |
| 61 | + except OSError: |
| 62 | + print("SPI init failed. Check config.txt and raspi-config.") |
| 63 | + self.led_init_state = 0 |
| 64 | + |
| 65 | + def check_spi_state(self): |
| 66 | + return self.led_init_state |
| 67 | + |
| 68 | + def show(self): |
| 69 | + d = numpy.array(self.led_color).ravel() |
| 70 | + tx = numpy.zeros(len(d) * 8, dtype=numpy.uint8) |
| 71 | + for ibit in range(8): |
| 72 | + tx[7 - ibit::8] = ((d >> ibit) & 1) * 0x78 + 0x80 |
| 73 | + if self.led_init_state: |
| 74 | + speed = int(8 / (1.25e-6 if self.bus == 0 else 1.0e-6)) |
| 75 | + self.spi.xfer(tx.tolist(), speed) |
| 76 | + |
| 77 | + def led_close(self): |
| 78 | + self.set_all_led_color(0, 0, 0) |
| 79 | + self.spi.close() |
| 80 | + |
| 81 | + def rhythm_wave(self, frame): |
| 82 | + colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255)] |
| 83 | + count = self.led_count |
| 84 | + base = (frame // 20) % len(colors) |
| 85 | + self.set_all_led_color(0, 0, 0) |
| 86 | + for i in range(count): |
| 87 | + phase = (frame + i * 6) % (count * 12) |
| 88 | + bright = phase / (count * 6) if phase < count * 6 else (count * 12 - phase) / (count * 6) |
| 89 | + color = colors[(base + i) % len(colors)] |
| 90 | + scaled = [int(c * bright) for c in color] |
| 91 | + self.set_ledpixel(i, *scaled) |
| 92 | + self.show() |
0 commit comments