alexei.dolgolyov
2b953e2e3e
Some checks failed
Validate / validate (push) Failing after 8s
Allow LEDs to cover only a fraction of each screen edge via draggable span bars in the calibration UI. Per-edge start/end (0.0-1.0) values control which portion of the screen border is sampled for LED colors. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
420 lines
15 KiB
Python
420 lines
15 KiB
Python
"""Calibration system for mapping screen pixels to LED positions."""
|
||
|
||
from dataclasses import dataclass, field
|
||
from typing import Dict, List, Literal, Tuple
|
||
|
||
from wled_controller.core.screen_capture import (
|
||
BorderPixels,
|
||
get_edge_segments,
|
||
calculate_average_color,
|
||
calculate_median_color,
|
||
calculate_dominant_color,
|
||
)
|
||
from wled_controller.utils import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
# Edge traversal order for each (start_position, layout) combination.
|
||
# Determines which edge comes first when walking around the screen.
|
||
EDGE_ORDER: Dict[Tuple[str, str], List[str]] = {
|
||
("bottom_left", "clockwise"): ["left", "top", "right", "bottom"],
|
||
("bottom_left", "counterclockwise"): ["bottom", "right", "top", "left"],
|
||
("bottom_right", "clockwise"): ["bottom", "left", "top", "right"],
|
||
("bottom_right", "counterclockwise"): ["right", "top", "left", "bottom"],
|
||
("top_left", "clockwise"): ["top", "right", "bottom", "left"],
|
||
("top_left", "counterclockwise"): ["left", "bottom", "right", "top"],
|
||
("top_right", "clockwise"): ["right", "bottom", "left", "top"],
|
||
("top_right", "counterclockwise"): ["top", "left", "bottom", "right"],
|
||
}
|
||
|
||
# Whether LEDs are reversed on each edge for each (start_position, layout) combination.
|
||
EDGE_REVERSE: Dict[Tuple[str, str], Dict[str, bool]] = {
|
||
("bottom_left", "clockwise"): {"left": True, "top": False, "right": False, "bottom": True},
|
||
("bottom_left", "counterclockwise"): {"bottom": False, "right": True, "top": True, "left": False},
|
||
("bottom_right", "clockwise"): {"bottom": True, "left": True, "top": False, "right": False},
|
||
("bottom_right", "counterclockwise"): {"right": True, "top": True, "left": False, "bottom": False},
|
||
("top_left", "clockwise"): {"top": False, "right": False, "bottom": True, "left": True},
|
||
("top_left", "counterclockwise"): {"left": False, "bottom": False, "right": True, "top": True},
|
||
("top_right", "clockwise"): {"right": False, "bottom": True, "left": True, "top": False},
|
||
("top_right", "counterclockwise"): {"top": True, "left": False, "bottom": False, "right": True},
|
||
}
|
||
|
||
|
||
@dataclass
|
||
class CalibrationSegment:
|
||
"""Configuration for one segment of the LED strip."""
|
||
|
||
edge: Literal["top", "right", "bottom", "left"]
|
||
led_start: int
|
||
led_count: int
|
||
reverse: bool = False
|
||
|
||
|
||
@dataclass
|
||
class CalibrationConfig:
|
||
"""Complete calibration configuration.
|
||
|
||
Stores only the core parameters. Segments (with led_start, reverse, edge order)
|
||
are derived at runtime via the `segments` property.
|
||
"""
|
||
|
||
layout: Literal["clockwise", "counterclockwise"]
|
||
start_position: Literal["top_left", "top_right", "bottom_left", "bottom_right"]
|
||
offset: int = 0
|
||
leds_top: int = 0
|
||
leds_right: int = 0
|
||
leds_bottom: int = 0
|
||
leds_left: int = 0
|
||
# Per-edge span: fraction of screen side covered by LEDs (0.0–1.0)
|
||
span_top_start: float = 0.0
|
||
span_top_end: float = 1.0
|
||
span_right_start: float = 0.0
|
||
span_right_end: float = 1.0
|
||
span_bottom_start: float = 0.0
|
||
span_bottom_end: float = 1.0
|
||
span_left_start: float = 0.0
|
||
span_left_end: float = 1.0
|
||
|
||
def build_segments(self) -> List[CalibrationSegment]:
|
||
"""Derive segment list from core parameters."""
|
||
key = (self.start_position, self.layout)
|
||
edge_order = EDGE_ORDER.get(key, ["bottom", "right", "top", "left"])
|
||
reverse_map = EDGE_REVERSE.get(key, {})
|
||
|
||
led_counts = {
|
||
"top": self.leds_top,
|
||
"right": self.leds_right,
|
||
"bottom": self.leds_bottom,
|
||
"left": self.leds_left,
|
||
}
|
||
|
||
segments = []
|
||
led_start = 0
|
||
for edge in edge_order:
|
||
count = led_counts[edge]
|
||
if count > 0:
|
||
segments.append(CalibrationSegment(
|
||
edge=edge,
|
||
led_start=led_start,
|
||
led_count=count,
|
||
reverse=reverse_map.get(edge, False),
|
||
))
|
||
led_start += count
|
||
|
||
return segments
|
||
|
||
@property
|
||
def segments(self) -> List[CalibrationSegment]:
|
||
"""Get derived segment list."""
|
||
return self.build_segments()
|
||
|
||
def get_edge_span(self, edge: str) -> tuple[float, float]:
|
||
"""Get span (start, end) for a given edge."""
|
||
return (
|
||
getattr(self, f"span_{edge}_start", 0.0),
|
||
getattr(self, f"span_{edge}_end", 1.0),
|
||
)
|
||
|
||
def validate(self) -> bool:
|
||
"""Validate calibration configuration.
|
||
|
||
Returns:
|
||
True if configuration is valid
|
||
|
||
Raises:
|
||
ValueError: If configuration is invalid
|
||
"""
|
||
total = self.get_total_leds()
|
||
if total <= 0:
|
||
raise ValueError("Calibration must have at least one LED")
|
||
|
||
for edge, count in [("top", self.leds_top), ("right", self.leds_right),
|
||
("bottom", self.leds_bottom), ("left", self.leds_left)]:
|
||
if count < 0:
|
||
raise ValueError(f"LED count for {edge} must be non-negative, got {count}")
|
||
|
||
for edge in ["top", "right", "bottom", "left"]:
|
||
start, end = self.get_edge_span(edge)
|
||
if not (0.0 <= start <= 1.0) or not (0.0 <= end <= 1.0):
|
||
raise ValueError(f"Span for {edge} must be in [0.0, 1.0], got ({start}, {end})")
|
||
if end <= start:
|
||
raise ValueError(f"Span end must be greater than start for {edge}, got ({start}, {end})")
|
||
|
||
return True
|
||
|
||
def get_total_leds(self) -> int:
|
||
"""Get total number of LEDs across all edges."""
|
||
return self.leds_top + self.leds_right + self.leds_bottom + self.leds_left
|
||
|
||
def get_segment_for_edge(self, edge: str) -> CalibrationSegment | None:
|
||
"""Get segment configuration for a specific edge."""
|
||
for seg in self.segments:
|
||
if seg.edge == edge:
|
||
return seg
|
||
return None
|
||
|
||
|
||
class PixelMapper:
|
||
"""Maps screen border pixels to LED colors based on calibration."""
|
||
|
||
def __init__(
|
||
self,
|
||
calibration: CalibrationConfig,
|
||
interpolation_mode: Literal["average", "median", "dominant"] = "average",
|
||
):
|
||
"""Initialize pixel mapper.
|
||
|
||
Args:
|
||
calibration: Calibration configuration
|
||
interpolation_mode: Color calculation mode
|
||
"""
|
||
self.calibration = calibration
|
||
self.interpolation_mode = interpolation_mode
|
||
|
||
# Validate calibration
|
||
self.calibration.validate()
|
||
|
||
# Select color calculation function
|
||
if interpolation_mode == "average":
|
||
self._calc_color = calculate_average_color
|
||
elif interpolation_mode == "median":
|
||
self._calc_color = calculate_median_color
|
||
elif interpolation_mode == "dominant":
|
||
self._calc_color = calculate_dominant_color
|
||
else:
|
||
raise ValueError(f"Invalid interpolation mode: {interpolation_mode}")
|
||
|
||
logger.info(
|
||
f"Initialized pixel mapper with {self.calibration.get_total_leds()} LEDs "
|
||
f"using {interpolation_mode} interpolation"
|
||
)
|
||
|
||
def map_border_to_leds(
|
||
self,
|
||
border_pixels: BorderPixels
|
||
) -> List[Tuple[int, int, int]]:
|
||
"""Map screen border pixels to LED colors.
|
||
|
||
Args:
|
||
border_pixels: Extracted border pixels from screen
|
||
|
||
Returns:
|
||
List of (R, G, B) tuples for each LED
|
||
|
||
Raises:
|
||
ValueError: If border pixels don't match calibration
|
||
"""
|
||
total_leds = self.calibration.get_total_leds()
|
||
led_colors = [(0, 0, 0)] * total_leds
|
||
|
||
# Process each edge
|
||
for edge_name in ["top", "right", "bottom", "left"]:
|
||
segment = self.calibration.get_segment_for_edge(edge_name)
|
||
|
||
if not segment:
|
||
# This edge is not configured
|
||
continue
|
||
|
||
# Get pixels for this edge
|
||
if edge_name == "top":
|
||
edge_pixels = border_pixels.top
|
||
elif edge_name == "right":
|
||
edge_pixels = border_pixels.right
|
||
elif edge_name == "bottom":
|
||
edge_pixels = border_pixels.bottom
|
||
else: # left
|
||
edge_pixels = border_pixels.left
|
||
|
||
# Slice to span region if not full coverage
|
||
span_start, span_end = self.calibration.get_edge_span(edge_name)
|
||
if span_start > 0.0 or span_end < 1.0:
|
||
if edge_name in ("top", "bottom"):
|
||
total_w = edge_pixels.shape[1]
|
||
s = int(span_start * total_w)
|
||
e = int(span_end * total_w)
|
||
edge_pixels = edge_pixels[:, s:e, :]
|
||
else:
|
||
total_h = edge_pixels.shape[0]
|
||
s = int(span_start * total_h)
|
||
e = int(span_end * total_h)
|
||
edge_pixels = edge_pixels[s:e, :, :]
|
||
|
||
# Divide edge into segments matching LED count
|
||
try:
|
||
pixel_segments = get_edge_segments(
|
||
edge_pixels,
|
||
segment.led_count,
|
||
edge_name
|
||
)
|
||
except ValueError as e:
|
||
logger.error(f"Failed to segment {edge_name} edge: {e}")
|
||
raise
|
||
|
||
# Calculate LED indices for this segment
|
||
led_indices = list(range(segment.led_start, segment.led_start + segment.led_count))
|
||
|
||
# Reverse if needed
|
||
if segment.reverse:
|
||
led_indices = list(reversed(led_indices))
|
||
|
||
# Map pixel segments to LEDs
|
||
for led_idx, pixel_segment in zip(led_indices, pixel_segments):
|
||
color = self._calc_color(pixel_segment)
|
||
led_colors[led_idx] = color
|
||
|
||
# Apply physical LED offset by rotating the array
|
||
# Offset = number of LEDs from LED 0 to the start corner
|
||
# Physical LED[i] should get calibration color[(i - offset) % total]
|
||
offset = self.calibration.offset % total_leds if total_leds > 0 else 0
|
||
if offset > 0:
|
||
led_colors = led_colors[total_leds - offset:] + led_colors[:total_leds - offset]
|
||
|
||
logger.debug(f"Mapped border pixels to {total_leds} LED colors (offset={offset})")
|
||
return led_colors
|
||
|
||
def test_calibration(self, edge: str, color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
|
||
"""Generate test pattern to light up specific edge.
|
||
|
||
Useful for verifying calibration configuration.
|
||
|
||
Args:
|
||
edge: Edge to light up (top, right, bottom, left)
|
||
color: RGB color to use
|
||
|
||
Returns:
|
||
List of LED colors with only the specified edge lit
|
||
|
||
Raises:
|
||
ValueError: If edge is not in calibration
|
||
"""
|
||
segment = self.calibration.get_segment_for_edge(edge)
|
||
if not segment:
|
||
raise ValueError(f"Edge '{edge}' not found in calibration")
|
||
|
||
total_leds = self.calibration.get_total_leds()
|
||
led_colors = [(0, 0, 0)] * total_leds
|
||
|
||
# Light up the specified edge
|
||
led_indices = range(segment.led_start, segment.led_start + segment.led_count)
|
||
for led_idx in led_indices:
|
||
led_colors[led_idx] = color
|
||
|
||
logger.info(f"Generated test pattern for {edge} edge with color {color}")
|
||
return led_colors
|
||
|
||
|
||
def create_default_calibration(led_count: int) -> CalibrationConfig:
|
||
"""Create a default calibration for a rectangular screen.
|
||
|
||
Assumes LEDs are evenly distributed around the screen edges in clockwise order
|
||
starting from bottom-left.
|
||
|
||
Args:
|
||
led_count: Total number of LEDs
|
||
|
||
Returns:
|
||
Default calibration configuration
|
||
"""
|
||
if led_count < 4:
|
||
raise ValueError("Need at least 4 LEDs for default calibration")
|
||
|
||
# Distribute LEDs evenly across 4 edges
|
||
leds_per_edge = led_count // 4
|
||
remainder = led_count % 4
|
||
|
||
# Distribute remainder to longer edges (bottom and top)
|
||
bottom_count = leds_per_edge + (1 if remainder > 0 else 0)
|
||
right_count = leds_per_edge
|
||
top_count = leds_per_edge + (1 if remainder > 1 else 0)
|
||
left_count = leds_per_edge + (1 if remainder > 2 else 0)
|
||
|
||
config = CalibrationConfig(
|
||
layout="clockwise",
|
||
start_position="bottom_left",
|
||
leds_bottom=bottom_count,
|
||
leds_right=right_count,
|
||
leds_top=top_count,
|
||
leds_left=left_count,
|
||
)
|
||
|
||
logger.info(
|
||
f"Created default calibration for {led_count} LEDs: "
|
||
f"bottom={bottom_count}, right={right_count}, "
|
||
f"top={top_count}, left={left_count}"
|
||
)
|
||
|
||
return config
|
||
|
||
|
||
def calibration_from_dict(data: dict) -> CalibrationConfig:
|
||
"""Create calibration configuration from dictionary.
|
||
|
||
Supports both new format (leds_top/right/bottom/left) and legacy format
|
||
(segments list) for backward compatibility.
|
||
|
||
Args:
|
||
data: Dictionary with calibration data
|
||
|
||
Returns:
|
||
CalibrationConfig instance
|
||
|
||
Raises:
|
||
ValueError: If data is invalid
|
||
"""
|
||
try:
|
||
config = CalibrationConfig(
|
||
layout=data["layout"],
|
||
start_position=data["start_position"],
|
||
offset=data.get("offset", 0),
|
||
leds_top=data.get("leds_top", 0),
|
||
leds_right=data.get("leds_right", 0),
|
||
leds_bottom=data.get("leds_bottom", 0),
|
||
leds_left=data.get("leds_left", 0),
|
||
span_top_start=data.get("span_top_start", 0.0),
|
||
span_top_end=data.get("span_top_end", 1.0),
|
||
span_right_start=data.get("span_right_start", 0.0),
|
||
span_right_end=data.get("span_right_end", 1.0),
|
||
span_bottom_start=data.get("span_bottom_start", 0.0),
|
||
span_bottom_end=data.get("span_bottom_end", 1.0),
|
||
span_left_start=data.get("span_left_start", 0.0),
|
||
span_left_end=data.get("span_left_end", 1.0),
|
||
)
|
||
|
||
config.validate()
|
||
return config
|
||
|
||
except KeyError as e:
|
||
raise ValueError(f"Missing required calibration field: {e}")
|
||
except Exception as e:
|
||
if isinstance(e, ValueError):
|
||
raise
|
||
raise ValueError(f"Invalid calibration data: {e}")
|
||
|
||
|
||
def calibration_to_dict(config: CalibrationConfig) -> dict:
|
||
"""Convert calibration configuration to dictionary.
|
||
|
||
Args:
|
||
config: Calibration configuration
|
||
|
||
Returns:
|
||
Dictionary representation
|
||
"""
|
||
result = {
|
||
"layout": config.layout,
|
||
"start_position": config.start_position,
|
||
"offset": config.offset,
|
||
"leds_top": config.leds_top,
|
||
"leds_right": config.leds_right,
|
||
"leds_bottom": config.leds_bottom,
|
||
"leds_left": config.leds_left,
|
||
}
|
||
# Include span fields only when not default (full coverage)
|
||
for edge in ["top", "right", "bottom", "left"]:
|
||
start = getattr(config, f"span_{edge}_start")
|
||
end = getattr(config, f"span_{edge}_end")
|
||
if start != 0.0 or end != 1.0:
|
||
result[f"span_{edge}_start"] = start
|
||
result[f"span_{edge}_end"] = end
|
||
return result
|