From 4ccfcef667816792945c55f5d474f732f9b3f685 Mon Sep 17 00:00:00 2001
From: Antoine Beyeler <49431240+abey79@users.noreply.github.com>
Date: Mon, 16 Dec 2024 10:09:19 +0100
Subject: [PATCH] Add utility to `rr.components.Color` to generate colors from
 any string (and use it in the air traffic data example) (#8458)

### What

It's some time nice to log some color information in multiple entities
to make it easier to relate them visually. This PR adds a
`rr.components.Color.from_str()` utility that does exactly that:
generate a nice color randomly picked up based on the provided string.

This PR also updates the air traffic data example so the barometric
traces have matching colors with the map data.

---------

Co-authored-by: Clement Rey <cr.rey.clement@gmail.com>
---
 .../air_traffic_data/air_traffic_data.py      | 27 ++++++++++++------
 rerun_py/rerun_sdk/rerun/components/color.py  |  3 +-
 .../rerun_sdk/rerun/components/color_ext.py   | 28 +++++++++++++++++++
 3 files changed, 49 insertions(+), 9 deletions(-)
 create mode 100644 rerun_py/rerun_sdk/rerun/components/color_ext.py

diff --git a/examples/python/air_traffic_data/air_traffic_data.py b/examples/python/air_traffic_data/air_traffic_data.py
index 911c4423e69d..2e2352e0ef0c 100644
--- a/examples/python/air_traffic_data/air_traffic_data.py
+++ b/examples/python/air_traffic_data/air_traffic_data.py
@@ -239,6 +239,7 @@ def process_measurement(self, measurement: Measurement) -> None:
             )
 
         entity_path = f"aircraft/{measurement.icao_id}"
+        color = rr.components.Color.from_string(entity_path)
 
         if (
             measurement.latitude is not None
@@ -247,13 +248,16 @@ def process_measurement(self, measurement: Measurement) -> None:
         ):
             rr.log(
                 entity_path,
-                rr.Points3D([
-                    self._proj.transform(
-                        measurement.longitude,
-                        measurement.latitude,
-                        measurement.barometric_altitude,
-                    )
-                ]),
+                rr.Points3D(
+                    [
+                        self._proj.transform(
+                            measurement.longitude,
+                            measurement.latitude,
+                            measurement.barometric_altitude,
+                        ),
+                    ],
+                    colors=color,
+                ),
                 rr.GeoPoints(lat_lon=[measurement.latitude, measurement.longitude]),
             )
 
@@ -264,6 +268,7 @@ def process_measurement(self, measurement: Measurement) -> None:
             rr.log(
                 entity_path + "/barometric_altitude",
                 rr.Scalar(measurement.barometric_altitude),
+                rr.SeriesLine(color=color),
             )
 
     def flush(self) -> None:
@@ -310,7 +315,13 @@ def log_position_and_altitude(self, df: polars.DataFrame, icao_id: str) -> None:
             return
 
         if icao_id not in self._position_indicators:
-            rr.log(entity_path, [rr.archetypes.Points3D.indicator(), rr.archetypes.GeoPoints.indicator()], static=True)
+            color = rr.components.Color.from_string(entity_path)
+            rr.log(
+                entity_path,
+                [rr.archetypes.Points3D.indicator(), rr.archetypes.GeoPoints.indicator(), color],
+                static=True,
+            )
+            rr.log(entity_path + "/barometric_altitude", [rr.archetypes.SeriesLine.indicator(), color], static=True)
             self._position_indicators.add(icao_id)
 
         timestamps = rr.TimeSecondsColumn("unix_time", df["timestamp"].to_numpy())
diff --git a/rerun_py/rerun_sdk/rerun/components/color.py b/rerun_py/rerun_sdk/rerun/components/color.py
index 5d72204fde68..122a0a9cfd7f 100644
--- a/rerun_py/rerun_sdk/rerun/components/color.py
+++ b/rerun_py/rerun_sdk/rerun/components/color.py
@@ -11,11 +11,12 @@
     ComponentDescriptor,
     ComponentMixin,
 )
+from .color_ext import ColorExt
 
 __all__ = ["Color", "ColorBatch"]
 
 
-class Color(datatypes.Rgba32, ComponentMixin):
+class Color(ColorExt, datatypes.Rgba32, ComponentMixin):
     """
     **Component**: An RGBA color with unmultiplied/separate alpha, in sRGB gamma space with linear alpha.
 
diff --git a/rerun_py/rerun_sdk/rerun/components/color_ext.py b/rerun_py/rerun_sdk/rerun/components/color_ext.py
new file mode 100644
index 000000000000..b06d2f7c9624
--- /dev/null
+++ b/rerun_py/rerun_sdk/rerun/components/color_ext.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import colorsys
+import math
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from . import Color
+
+_GOLDEN_RATIO = (math.sqrt(5.0) - 1.0) / 2.0
+
+
+class ColorExt:
+    """Extension for [Color][rerun.components.Color]."""
+
+    @staticmethod
+    def from_string(s: str) -> Color:
+        """
+        Generate a random yet deterministic color based on a string.
+
+        The color is guaranteed to be identical for the same input string.
+        """
+
+        from . import Color
+
+        # adapted from egui::PlotUi
+        hue = (hash(s) & 0xFFFF) / 2**16 * _GOLDEN_RATIO
+        return Color([round(comp * 255) for comp in colorsys.hsv_to_rgb(hue, 0.85, 0.5)])