Add backend logic for bounding box plots (#5250)

extension/src/cli/dvc/contract.ts

@@ -233,10 +233,16 @@ export type TemplatePlotOutput = {
   type: PlotsType
 }
 
+export type BoundingBox = {
+  box: { left: number; right: number; top: number; bottom: number }
+  score: number
+}
+
 export type ImagePlotOutput = {
   revisions: string[]
   type: PlotsType
   url: string
+  boxes?: { [label: string]: BoundingBox[] }
 }
 
 export type PlotOutput = TemplatePlotOutput | ImagePlotOutput

extension/src/common/colors.ts

@@ -0,0 +1,41 @@
+const colorsList = [
+  '#945dd6',
+  '#13adc7',
+  '#f46837',
+  '#48bb78',
+  '#4299e1',
+  '#ed8936',
+  '#f56565'
+] as const
+
+export type Color = (typeof colorsList)[number]
+
+export const copyOriginalColors = (): Color[] => [...colorsList]
+
+const boundingBoxColorsList = [
+  '#ff3838',
+  '#ff9d97',
+  '#ff701f',
+  '#ffb21d',
+  '#cfd231',
+  '#48f90a',
+  '#92cc17',
+  '#3ddb86',
+  '#1a9334',
+  '#00d4bb',
+  '#2c99a8',
+  '#00c2ff',
+  '#344593',
+  '#6473ff',
+  '#0018ec',
+  '#8438ff',
+  '#520085',
+  '#cb38ff',
+  '#ff95c8',
+  '#ff37c7'
+] as const
+
+export type BoundingBoxColor = (typeof boundingBoxColorsList)[number]
+
+export const getBoundingBoxColor = (ind: number): BoundingBoxColor =>
+  boundingBoxColorsList[ind % boundingBoxColorsList.length]

extension/src/experiments/index.ts

@@ -27,7 +27,6 @@ import {
   pickFilterToAdd,
   pickFiltersToRemove
 } from './model/filterBy/quickPick'
-import { Color } from './model/status/colors'
 import { MAX_SELECTED_EXPERIMENTS, UNSELECTED } from './model/status'
 import { starredSort } from './model/sortBy/constants'
 import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
@@ -39,6 +38,7 @@ import { Experiment, ColumnType, TableData, Column } from './webview/contract'
 import { WebviewMessages } from './webview/messages'
 import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
+import { Color } from '../common/colors'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
 import {

extension/src/experiments/model/index.test.ts

@@ -2,7 +2,7 @@
 import { join } from 'path'
 import { commands } from 'vscode'
 import { ExperimentsModel } from '.'
-import { copyOriginalColors } from './status/colors'
+import { copyOriginalColors } from '../../common/colors'
 import gitLogFixture from '../../test/fixtures/expShow/base/gitLog'
 import rowOrderFixture from '../../test/fixtures/expShow/base/rowOrder'
 import outputFixture from '../../test/fixtures/expShow/base/output'

extension/src/experiments/model/index.ts

@@ -16,9 +16,9 @@ import {
   collectSelectedColors,
   collectStartedRunningExperiments
 } from './status/collect'
-import { Color, copyOriginalColors } from './status/colors'
 import { canSelect, ColoredStatus, UNSELECTED } from './status'
 import { collectFlatExperimentParams } from './modify/collect'
+import { Color, copyOriginalColors } from '../../common/colors'
 import {
   Commit,
   Experiment,

extension/src/experiments/model/status/collect.test.ts

@@ -1,6 +1,6 @@
 import { UNSELECTED } from '.'
 import { collectColoredStatus } from './collect'
-import { copyOriginalColors } from './colors'
+import { copyOriginalColors } from '../../../common/colors'
 import { Experiment } from '../../webview/contract'
 import {
   ExecutorStatus,

extension/src/experiments/model/status/collect.ts

@@ -5,7 +5,7 @@ import {
   tooManySelected,
   UNSELECTED
 } from '.'
-import { Color, copyOriginalColors } from './colors'
+import { Color, copyOriginalColors } from '../../../common/colors'
 import { hasKey } from '../../../util/object'
 import { Experiment, isQueued, RunningExperiment } from '../../webview/contract'
 import { definedAndNonEmpty, reorderListSubset } from '../../../util/array'

extension/src/experiments/model/status/index.test.ts

@@ -1,5 +1,5 @@
 import { canSelect, limitToMaxSelected } from '.'
-import { copyOriginalColors } from './colors'
+import { copyOriginalColors } from '../../../common/colors'
 import { Experiment } from '../../webview/contract'
 import { ExecutorStatus } from '../../../cli/dvc/contract'
 

extension/src/experiments/model/status/index.ts

@@ -1,4 +1,4 @@
-import { Color } from './colors'
+import { Color } from '../../../common/colors'
 import { Experiment, isRunning } from '../../webview/contract'
 
 export const MAX_SELECTED_EXPERIMENTS = 7

extension/src/persistence/constants.ts

@@ -20,6 +20,7 @@ export enum PersistenceKey {
   PLOT_SELECTED_METRICS = 'plotSelectedMetrics:',
   PLOTS_SMOOTH_PLOT_VALUES = 'plotSmoothPlotValues:',
   PLOTS_COMPARISON_MULTI_PLOT_VALUES = 'plotComparisonMultiPlotValues:',
+  PLOTS_COMPARISON_CLASSES_SELECTED = 'plotComparisonClassesSelected:',
   PLOT_TEMPLATE_ORDER = 'plotTemplateOrder:',
   SHOW_ONLY_CHANGED = 'columnsShowOnlyChanged:'
 }

extension/src/plots/model/collect.test.ts

@@ -114,13 +114,14 @@ describe('collectData', () => {
     ])
 
     const heatmapPlot = join('plots', 'heatmap.png')
+    const boundingBoxPlot = join('plots', 'bounding_boxes.png')
 
     expect(Object.keys(comparisonData.main)).toStrictEqual([
       join('plots', 'acc.png'),
       heatmapPlot,
       join('plots', 'loss.png'),
       join('plots', 'image'),
-      join('plots', 'bounding_boxes.png')
+      boundingBoxPlot
     ])
 
     const testBranchHeatmap = comparisonData['test-branch'][heatmapPlot]