fix: choose detection box colour by class name rather than index

getDetectionBoxColor() maps a raw class index to a colour using the COCO
ordering: 0 is person (blue), 1-8 vehicles (green), 14-23 animals (orange).
That is only meaningful for a model that actually is COCO. A custom model -
a firearm detector whose class 0 is "gun", say - got its weapons drawn in
the blue reserved for people.

Resolve the colour through the ObjectClasses instance's own class list
instead, so the decision is made on the class name. COCO is unaffected: a
test asserts the new lookup agrees with the legacy index mapping across all
80 classes. A custom model now gets red for classes we have no opinion on,
and still gets blue for "person" wherever that class happens to sit in its
ordering.

The static index-based functions are kept for compatibility but documented
as COCO-only.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ThMq92LnrBbU2y9astkCDJ

Author: connortechnology

src/zm_netint_yolo.cpp

@@ -436,7 +436,7 @@ int Quadra_Yolo::draw_roi_box_in_place(
 
   // Use provided color, or fall back to default class-based color
   if (box_color == 0) {
-    box_color = ObjectClasses::getDetectionBoxColor(roi_extra.cls);
+    box_color = object_classes_.boxColorFor(roi_extra.cls);
   }
   Image in_image(inframe);
 
@@ -463,7 +463,7 @@ int Quadra_Yolo::draw_roi_box(
     snprintf(color_buf, sizeof(color_buf), "0x%06X", box_color);
     color = color_buf;
   } else {
-    color = ObjectClasses::getDetectionColorString(roi_extra.cls);
+    color = object_classes_.colorStringFor(roi_extra.cls);
   }
 
   for (int i=0; i<line_width; i++) {

src/zm_object_classes.cpp

@@ -21,6 +21,7 @@
 #include "zm_logger.h"
 
 #include <fstream>
+#include <set>
 
 // Default COCO dataset class names (80 classes)
 const std::vector<std::string> ObjectClasses::kCocoClassNames = {
@@ -94,6 +95,41 @@ const std::string& ObjectClasses::getClassName(int class_id) const {
   return kUnknownClass;
 }
 
+namespace {
+
+// The COCO groupings the index-based mapping below encodes, by name so they
+// survive a model that numbers its classes differently.
+const std::set<std::string> kVehicleClasses = {
+  "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat"
+};
+const std::set<std::string> kAnimalClasses = {
+  "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe"
+};
+
+}  // namespace
+
+Rgb ObjectClasses::boxColorForName(const std::string &class_name) {
+  if (class_name == "person") return kRGBBlue;
+  if (kVehicleClasses.count(class_name)) return kRGBGreen;
+  if (kAnimalClasses.count(class_name)) return kRGBOrange;
+  return kRGBRed;
+}
+
+const char *ObjectClasses::colorStringForName(const std::string &class_name) {
+  if (class_name == "person") return "blue";
+  if (kVehicleClasses.count(class_name)) return "green";
+  if (kAnimalClasses.count(class_name)) return "orange";
+  return "red";
+}
+
+Rgb ObjectClasses::boxColorFor(int class_id) const {
+  return boxColorForName(getClassName(class_id));
+}
+
+const char *ObjectClasses::colorStringFor(int class_id) const {
+  return colorStringForName(getClassName(class_id));
+}
+
 Rgb ObjectClasses::getDetectionBoxColor(int class_id) {
   if (class_id == 0) {
     return kRGBBlue;    // Person

src/zm_object_classes.h

@@ -46,11 +46,28 @@ class ObjectClasses {
 
   // Get detection box color based on class ID.
   // Person (class 0) = Blue, Vehicles (1-8) = Green, Animals (14-23) = Orange, Others = Red
+  //
+  // These take a raw index and can therefore only assume the COCO ordering. A
+  // custom model's class 0 is not "person", so prefer the instance methods
+  // below, which resolve the index through this object's own class list.
   static Rgb getDetectionBoxColor(int class_id);
 
   // Get detection color name as string.
   static const char* getDetectionColorString(int class_id);
 
+  /* Colour for a detection box, chosen from the class *name* rather than its
+   * index. An index only means anything within the dataset it came from, so a
+   * firearm model whose class 0 is "gun" would otherwise be drawn in COCO's
+   * "person" blue. Resolving by name gives identical results for COCO and
+   * something sensible - red, the default - for classes we have no opinion on.
+   */
+  Rgb boxColorFor(int class_id) const;
+  const char *colorStringFor(int class_id) const;
+
+  // Name-based colour lookup, exposed so it can be tested directly.
+  static Rgb boxColorForName(const std::string &class_name);
+  static const char *colorStringForName(const std::string &class_name);
+
   // Access to underlying vector for iteration
   const std::vector<std::string>& getClassNames() const { return class_names_; }
 

tests/CMakeLists.txt

@@ -22,6 +22,7 @@ set(TEST_SOURCES
   zm_image.cpp
   zm_monitorstream.cpp
   zm_netint_lpr.cpp
+  zm_object_classes.cpp
   zm_onvif_renewal.cpp
   zm_onvif_wsse.cpp
   zm_pixformat.cpp

tests/zm_object_classes.cpp

@@ -0,0 +1,127 @@
+/*
+ * This file is part of the ZoneMinder Project. See AUTHORS file for Copyright information
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License as published by the
+ * Free Software Foundation; either version 2 of the License, or (at your
+ * option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "zm_catch2.h"
+
+#include <cstdio>
+#include <fstream>
+#include <string>
+
+#include "zm_object_classes.h"
+
+namespace {
+
+// Write a .names file beside a notional model so loadFromFile() finds it.
+std::string WriteNames(const std::string &stem, const std::string &contents) {
+  const std::string names = stem + ".names";
+  std::ofstream out(names);
+  out << contents;
+  out.close();
+  return stem + ".nb";  // what loadFromFile is handed
+}
+
+}  // namespace
+
+TEST_CASE("ObjectClasses::boxColorForName maps the COCO groupings", "[objectclasses]") {
+  // These are the groupings the old index-based mapping encoded. Resolving them
+  // by name has to give exactly the same answers for COCO.
+  SECTION("person is blue") {
+    REQUIRE(ObjectClasses::boxColorForName("person") == kRGBBlue);
+  }
+
+  SECTION("vehicles are green") {
+    for (const char *name : {"bicycle", "car", "motorcycle", "airplane",
+                             "bus", "train", "truck", "boat"}) {
+      REQUIRE(ObjectClasses::boxColorForName(name) == kRGBGreen);
+    }
+  }
+
+  SECTION("animals are orange") {
+    for (const char *name : {"bird", "cat", "dog", "horse", "sheep",
+                             "cow", "elephant", "bear", "zebra", "giraffe"}) {
+      REQUIRE(ObjectClasses::boxColorForName(name) == kRGBOrange);
+    }
+  }
+
+  SECTION("anything else is red") {
+    REQUIRE(ObjectClasses::boxColorForName("backpack") == kRGBRed);
+    REQUIRE(ObjectClasses::boxColorForName("gun") == kRGBRed);
+    REQUIRE(ObjectClasses::boxColorForName("") == kRGBRed);
+  }
+}
+
+TEST_CASE("ObjectClasses::boxColorFor agrees with the index mapping for COCO",
+          "[objectclasses]") {
+  // A default-constructed instance carries COCO, so the new name-based lookup
+  // must reproduce the legacy index-based one across every class. This is the
+  // regression guard: the refactor may not change existing behaviour.
+  ObjectClasses coco;
+  for (size_t i = 0; i < coco.size(); i++) {
+    const int id = static_cast<int>(i);
+    REQUIRE(coco.boxColorFor(id) == ObjectClasses::getDetectionBoxColor(id));
+    REQUIRE(std::string(coco.colorStringFor(id)) ==
+            std::string(ObjectClasses::getDetectionColorString(id)));
+  }
+}
+
+TEST_CASE("ObjectClasses::boxColorFor does not give a custom model COCO's colours",
+          "[objectclasses]") {
+  /* The bug this fixes: a firearm model's class 0 is "gun", but the index-based
+   * mapping only knows COCO, where 0 is "person" - so a weapon was drawn in the
+   * blue reserved for people.
+   */
+  const std::string model = WriteNames("/tmp/zm_objclasses_firearm", "gun\nrifle\nknife\n");
+
+  ObjectClasses classes;
+  REQUIRE(classes.loadFromFile(model));
+  REQUIRE(classes.size() == 3);
+  REQUIRE(classes.getClassName(0) == "gun");
+
+  // The legacy index lookup still claims blue for class 0 ...
+  REQUIRE(ObjectClasses::getDetectionBoxColor(0) == kRGBBlue);
+  // ... but resolving through this model's own names does not.
+  REQUIRE(classes.boxColorFor(0) == kRGBRed);
+  REQUIRE(classes.boxColorFor(1) == kRGBRed);
+  REQUIRE(classes.boxColorFor(2) == kRGBRed);
+  REQUIRE(std::string(classes.colorStringFor(0)) == "red");
+
+  std::remove("/tmp/zm_objclasses_firearm.names");
+}
+
+TEST_CASE("ObjectClasses::boxColorFor still honours names a custom model shares with COCO",
+          "[objectclasses]") {
+  // A custom model that reorders COCO-ish names must colour by meaning, not by
+  // position: "person" is blue wherever it lands.
+  const std::string model = WriteNames("/tmp/zm_objclasses_mixed", "gun\nperson\ncar\n");
+
+  ObjectClasses classes;
+  REQUIRE(classes.loadFromFile(model));
+  REQUIRE(classes.boxColorFor(0) == kRGBRed);     // gun
+  REQUIRE(classes.boxColorFor(1) == kRGBBlue);    // person, at index 1 not 0
+  REQUIRE(classes.boxColorFor(2) == kRGBGreen);   // car, at index 2 not 2-of-COCO
+
+  std::remove("/tmp/zm_objclasses_mixed.names");
+}
+
+TEST_CASE("ObjectClasses::loadFromFile falls back to COCO when no .names exists",
+          "[objectclasses]") {
+  ObjectClasses classes;
+  REQUIRE_FALSE(classes.loadFromFile("/tmp/zm_objclasses_definitely_absent.nb"));
+  REQUIRE(classes.size() == 80);
+  REQUIRE(classes.getClassName(0) == "person");
+  REQUIRE(classes.boxColorFor(0) == kRGBBlue);
+}