Add support for custom legends with colors.

PiperOrigin-RevId: 865980035

Author: Google AI Edge

src/ui/src/components/visualizer/common/types.ts

@@ -912,6 +912,43 @@ export declare interface LegendConfig {
    * If set, rename the "outputs" label to this string.
    */
   renameOutputsTo?: string;
+
+  /**
+   * If set, add custom legend items to the legend panel.
+   *
+   * Each item has a color block and a label.
+   */
+  customLegendItems?: CustomLegendItem[];
+}
+
+/** A custom legend item. */
+export declare interface CustomLegendItem {
+  label: string;
+
+  /**
+   * The background color of the color block. Default is transparent.
+   */
+  backgroundColor?: string;
+
+  /**
+   * The border width of the color block. Default is 0.
+   */
+  borderWidth?: number;
+
+  /**
+   * The border color of the color block. Default is transparent.
+   */
+  borderColor?: string;
+
+  /**
+   * The border style of the color block. Default is solid.
+   */
+  borderStyle?: string;
+
+  /**
+   * The border radius of the color block. Default is 0.
+   */
+  borderRadius?: number;
 }
 
 /** Config for "View on node" menu. */

src/ui/src/components/visualizer/legends_panel.ng.html

@@ -73,6 +73,25 @@
         </tr>
       }
     </ng-container>
+    <!-- Custom legend items -->
+    @for (item of customLegendItems; track item) {
+      <tr>
+        <td>
+          <div class="key custom-key"
+              [style.background-color]="item.backgroundColor ?? 'transparent'"
+              [style.border-width.px]="item.borderWidth ?? 0"
+              [style.border-color]="item.borderColor ?? 'transparent'"
+              [style.border-style]="item.borderStyle ?? 'solid'"
+              [style.border-radius.px]="item.borderRadius ?? 0">
+          </div>
+        </td>
+        <td>
+          <div class="label">
+            {{item.label}}
+          </div>
+        </td>
+      </tr>
+    }
     @if (!hideShortcuts) {
       @if (showDivider()) {
         <tr><td colspan="2"><div class="divider"></div></td></tr>

src/ui/src/components/visualizer/legends_panel.scss

@@ -94,6 +94,12 @@ $block_height: 14px;
       border: 2px solid var(--me-outgoing-edge-color);
     }
 
+    .custom-key {
+      width: $block_height;
+      height: $block_height;
+      margin-left: auto;
+    }
+
     .shortcut {
       color: var(--me-on-surface-variant-color);
       font-size: 10px;

src/ui/src/components/visualizer/legends_panel.ts

@@ -29,6 +29,7 @@ import {
 import {MatIconModule} from '@angular/material/icon';
 
 import {AppService} from './app_service';
+import {CustomLegendItem} from './common/types';
 
 /**
  * The panel that shows the lengends.
@@ -148,4 +149,8 @@ export class LegendsPanel {
   get hideShortcuts(): boolean {
     return this.appService.config()?.legendConfig?.hideShortcuts ?? false;
   }
+
+  get customLegendItems(): CustomLegendItem[] {
+    return this.appService.config()?.legendConfig?.customLegendItems ?? [];
+  }
 }