Improve color detection for typical use cases

lib/pbio/include/pbio/color.h

@@ -118,7 +118,11 @@ void pbio_color_to_hsv(pbio_color_t color, pbio_color_hsv_t *hsv);
 void pbio_color_to_rgb(pbio_color_t color, pbio_color_rgb_t *rgb);
 void pbio_color_hsv_compress(const pbio_color_hsv_t *hsv, pbio_color_compressed_hsv_t *compressed);
 void pbio_color_hsv_expand(const pbio_color_compressed_hsv_t *compressed, pbio_color_hsv_t *hsv);
-int32_t pbio_color_get_bicone_squared_distance(const pbio_color_hsv_t *hsv_a, const pbio_color_hsv_t *hsv_b);
+
+typedef int32_t (*pbio_color_distance_func_t)(const pbio_color_hsv_t *hsv_a, const pbio_color_hsv_t *hsv_b);
+
+int32_t pbio_color_get_distance_bicone_squared(const pbio_color_hsv_t *hsv_a, const pbio_color_hsv_t *hsv_b);
+int32_t pbio_color_get_distance_saturation_heuristic(const pbio_color_hsv_t *hsv_a, const pbio_color_hsv_t *hsv_b);
 
 #endif // _PBIO_COLOR_H_
 

lib/pbio/src/color/util.c

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2022 The Pybricks Authors
 
+#include <assert.h>
 #include <pbio/color.h>
 #include <pbio/int_math.h>
 
@@ -13,7 +14,7 @@
  * @param [in]  hsv_b    The second HSV color.
  * @returns              Squared distance (0 to 400000000).
  */
-int32_t pbio_color_get_bicone_squared_distance(const pbio_color_hsv_t *hsv_a, const pbio_color_hsv_t *hsv_b) {
+int32_t pbio_color_get_distance_bicone_squared(const pbio_color_hsv_t *hsv_a, const pbio_color_hsv_t *hsv_b) {
 
     // Chroma (= radial coordinate in bicone) of a and b (0-10000).
     int32_t radius_a = pbio_color_hsv_get_v(hsv_a) * hsv_a->s;
@@ -36,3 +37,47 @@ int32_t pbio_color_get_bicone_squared_distance(const pbio_color_hsv_t *hsv_a, co
     // Squared Euclidean distance (0, 400000000)
     return delta_x * delta_x + delta_y * delta_y + delta_z * delta_z;
 }
+
+/**
+ * Gets distance measure between a HSV color (a) and a fully or zero saturated
+ * candidate color.
+ *
+ * @param [in]  measurement    The measured HSV color.
+ * @param [in]  candidate      The candidate HSV color (an idealized color or grayscale).
+ * @returns                    Heuristic distance.
+ */
+int32_t pbio_color_get_distance_saturation_heuristic(const pbio_color_hsv_t *measurement, const pbio_color_hsv_t *candidate) {
+
+    bool idealized_grayscale = candidate->s == 0 && candidate->h == 0;
+    bool idealized_color = candidate->s == 100 && candidate->v == 100;
+
+    // Calling code needs to ensure this.
+    assert(idealized_grayscale || idealized_color);
+
+    uint32_t hue_dist = pbio_int_math_abs(candidate->h - measurement->h);
+    if (hue_dist > 180) {
+        hue_dist = 360 - hue_dist;
+    }
+
+    uint32_t value_dist = pbio_int_math_abs(candidate->v - measurement->v);
+
+    const uint32_t penalty = 1000;
+
+    if (measurement->s <= 40 || measurement->v <= 1) {
+        // Measurement is unsaturated, so match to nearest grayscale; penalize color.
+        if (idealized_grayscale) {
+            // Match to nearest value.
+            return value_dist;
+        }
+        // Looking for grayscale, so disqualify color candidate.
+        return penalty + hue_dist;
+    } else {
+        // Measurement is saturated, so match to nearest full color; penalize grayscale.
+        if (idealized_color) {
+            // Match to nearest hue.
+            return hue_dist;
+        }
+        // Looking for color, so disqualify grayscale candidate.
+        return penalty + value_dist;
+    }
+}

lib/pbio/test/src/test_color.c

@@ -366,7 +366,7 @@ static void test_color_hsv_cost(void *env) {
     color_a.h = 0;
     color_a.s = 100;
     color_a.v = 100;
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_a), ==, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_a), ==, 0);
 
     // blacks with different saturations/hues should be the same
     color_a.h = 230;
@@ -376,7 +376,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.h = 23;
     color_b.s = 99;
     color_b.v = 0;
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), ==, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), ==, 0);
 
     // colors with different hues should be different when value>0 and saturation>0
     color_a.h = 230;
@@ -386,7 +386,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.h = 23;
     color_b.s = 99;
     color_b.v = 100;
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), >, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), >, 0);
 
     // grays with different hues should be the same
     color_a.h = 230;
@@ -396,7 +396,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.h = 23;
     color_b.s = 0;
     color_b.v = 50;
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), ==, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), ==, 0);
 
     // distance should be greater when saturation is greater
     color_a.h = 30;
@@ -407,7 +407,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 20;
     color_b.v = 70;
 
-    dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+    dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
 
     color_a.h = 30;
     color_a.s = 40;
@@ -417,7 +417,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 40;
     color_b.v = 70;
 
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), >, dist);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), >, dist);
 
     // resolve colors that are close
     color_a.h = 30;
@@ -428,7 +428,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 20;
     color_b.v = 70;
 
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), >, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), >, 0);
 
     color_a.h = 30;
     color_a.s = 20;
@@ -438,7 +438,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 25;
     color_b.v = 70;
 
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), >, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), >, 0);
 
     color_a.h = 30;
     color_a.s = 20;
@@ -448,7 +448,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 20;
     color_b.v = 75;
 
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), >, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), >, 0);
 
     // hues 360 and 0 should be the same
     color_a.h = 360;
@@ -458,7 +458,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.h = 0;
     color_b.s = 100;
     color_b.v = 100;
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), ==, 0);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), ==, 0);
 
     // distance between hues 359 and 1 should be smaller than hues 1 and 5
     color_a.h = 359;
@@ -468,7 +468,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.h = 1;
     color_b.s = 100;
     color_b.v = 100;
-    dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+    dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
 
     color_a.h = 1;
     color_a.s = 100;
@@ -478,7 +478,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 100;
     color_b.v = 100;
 
-    tt_want_int_op(pbio_color_get_bicone_squared_distance(&color_a, &color_b), >, dist);
+    tt_want_int_op(pbio_color_get_distance_bicone_squared(&color_a, &color_b), >, dist);
 
     // check distance is monotonous along several color paths. This should catch potential int overflows
     int prev_dist = 0;
@@ -495,7 +495,7 @@ static void test_color_hsv_cost(void *env) {
 
     while (color_a.s < 100) {
         color_a.s += 5;
-        dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+        dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
 
         if (dist <= prev_dist) {
             monotone = false;
@@ -520,7 +520,7 @@ static void test_color_hsv_cost(void *env) {
 
     while (color_a.v < 100) {
         color_a.v += 5;
-        dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+        dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
 
         if (dist <= prev_dist) {
             monotone = false;
@@ -545,7 +545,7 @@ static void test_color_hsv_cost(void *env) {
 
     while (color_a.v < 100) {
         color_a.v += 5;
-        dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+        dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
 
         if (dist <= prev_dist) {
             monotone = false;
@@ -573,7 +573,7 @@ static void test_color_hsv_cost(void *env) {
         color_a.h = i < 0 ? 180 : 0;
         color_a.v = 10000 / (200 - color_a.s); // constant lightness
 
-        dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+        dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
 
         if (dist <= prev_dist) {
             monotone = false;
@@ -592,7 +592,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 100;
     color_b.v = 100;
 
-    dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+    dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
     tt_want_int_op(dist, >, 390000000);
     tt_want_int_op(dist, <, 410000000);
 
@@ -604,7 +604,7 @@ static void test_color_hsv_cost(void *env) {
     color_b.s = 0;
     color_b.v = 100;
 
-    dist = pbio_color_get_bicone_squared_distance(&color_a, &color_b);
+    dist = pbio_color_get_distance_bicone_squared(&color_a, &color_b);
     tt_want_int_op(dist, >, 390000000);
     tt_want_int_op(dist, <, 410000000);
 }

pybricks/parameters/pb_type_color.c

@@ -62,7 +62,7 @@ const pb_type_Color_obj_t pb_Color_MAGENTA_obj = {
 
 const pb_type_Color_obj_t pb_Color_NONE_obj = {
     {&pb_type_Color},
-    .hsv = {0, 0, -40}
+    .hsv = {0, 0, 0}
 };
 
 const pb_type_Color_obj_t pb_Color_BLACK_obj = {

pybricks/pupdevices/pb_type_pupdevices_colordistancesensor.c

@@ -90,6 +90,10 @@ static void get_hsv_data(pupdevices_ColorDistanceSensor_obj_t *self, pbio_color_
     rgb.g = 1187 * raw[1] / 2048;
     rgb.b = 1187 * raw[2] / 2048;
     pb_color_map_rgb_to_hsv(&rgb, hsv);
+
+    // Approximately double low values to get similar results
+    // as with other sensors.
+    hsv->v = hsv->v * (200 - hsv->v) / 100;
 }
 
 // pybricks.pupdevices.ColorDistanceSensor.color

pybricks/pupdevices/pb_type_pupdevices_colorsensor.c

@@ -71,6 +71,14 @@ static void get_hsv_reflected(mp_obj_t self_in, pbio_color_hsv_t *hsv) {
         .b = data[2] == 1024 ? 255 : data[2] >> 2,
     };
     pb_color_map_rgb_to_hsv(&rgb, hsv);
+
+    // Approximately double saturation for low values to get similar results
+    // as other sensors.
+    hsv->s = hsv->s * (200 - hsv->s) / 100;
+
+    // Approximately +50% low values to get similar results
+    // as with other sensors.
+    hsv->v = hsv->v * (150 - hsv->v / 2) / 100;
 }
 
 // Helper for getting HSV with the light off, scale saturation and value to

pybricks/util_pb/pb_color_map.c

@@ -33,15 +33,13 @@ void pb_color_map_rgb_to_hsv(const pbio_color_rgb_t *rgb, pbio_color_hsv_t *hsv)
     pbio_color_rgb_to_hsv(rgb, hsv);
 
     // Slight shift for lower hues to make yellow somewhat more accurate
-    if (hsv->h < 40) {
-        uint8_t offset = ((hsv->h - 20) << 8) / 20;
-        int32_t scale = 200 - ((100 * (offset * offset)) >> 16);
-        hsv->h = hsv->h * scale / 100;
+    if (hsv->h >= 350) {
+        hsv->h = (350 + 2 * (hsv->h - 350)) % 360;
+    } else if (hsv->h < 40) {
+        hsv->h += 10;
+    } else if (hsv->h < 60) {
+        hsv->h = 50 + (hsv->h - 40) / 2;
     }
-
-    // Value and saturation correction
-    hsv->s = hsv->s * (200 - hsv->s) / 100;
-    hsv->v = hsv->v * (200 - hsv->v) / 100;
 }
 
 static const mp_rom_obj_tuple_t pb_color_map_default = {
@@ -75,11 +73,29 @@ mp_obj_t pb_color_map_get_color(mp_obj_t *color_map, pbio_color_hsv_t *hsv) {
     int32_t cost_now = INT32_MAX;
     int32_t cost_min = INT32_MAX;
 
+    pbio_color_distance_func_t distance_func = pbio_color_get_distance_saturation_heuristic;
+
+    // If user only provides fully saturated colors (hue, 100, 100) and/or fully
+    // desaturated colors (0, 0, value), use a simplified heuristic matcher for
+    // better default results that are distance independent. Otherwise use a
+    // bicone color distance measure.
+    for (size_t i = 0; i < n; i++) {
+        const pbio_color_hsv_t *candidate = pb_type_Color_get_hsv(colors[i]);
+
+        // Use bicone mapping if custom (realistic) colors provided.
+        bool idealized_grayscale = candidate->s == 0 && candidate->h == 0;
+        bool idealized_color = candidate->s == 100 && candidate->v == 100;
+        if (!idealized_grayscale && !idealized_color) {
+            distance_func = pbio_color_get_distance_bicone_squared;
+            break;
+        }
+    }
+
     // Compute cost for each candidate
     for (size_t i = 0; i < n; i++) {
 
         // Evaluate the cost function
-        cost_now = pbio_color_get_bicone_squared_distance(hsv, pb_type_Color_get_hsv(colors[i]));
+        cost_now = distance_func(hsv, pb_type_Color_get_hsv(colors[i]));
 
         // If cost is less than before, update the minimum and the match
         if (cost_now < cost_min) {