Add backend logic for bounding box plots

extension/src/cli/dvc/contract.ts

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

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]
@@ -131,6 +132,18 @@ describe('collectData', () => {
         sameContents(revisions, ['test-branch'])
       )
     ])
+
+    const workspaceBoundingBoxData = comparisonData.workspace[boundingBoxPlot]
+
+    expect(workspaceBoundingBoxData[0]).toMatchObject({
+      boundingBoxes: {
+        car: [
+          { bottom: 350, left: 150, right: 180, top: 320 },
+          { bottom: 340, left: 200, right: 230, top: 310 }
+        ],
+        'traffic light': [{ bottom: 210, left: 120, right: 195, top: 120 }]
+      }
+    })
   })
 })
 

extension/src/plots/model/collect.ts

@@ -14,7 +14,11 @@ import {
   CustomPlotData,
   CustomPlotValues,
   ComparisonRevisionData,
-  ComparisonPlotImg
+  ComparisonPlotImg,
+  ComparisonClassDetails,
+  ComparisonPlotClasses,
+  ComparisonClassesSelected,
+  ComparisonPlotClass
 } from '../webview/contract'
 import {
   AnchorDefinitions,
@@ -23,7 +27,9 @@ import {
   PlotOutput,
   PlotsOutput,
   PlotsType,
-  TemplatePlotOutput
+  TemplatePlotOutput,
+  ImagePlotOutput,
+  BoundingBox
 } from '../../cli/dvc/contract'
 import { splitColumnPath } from '../../experiments/columns/paths'
 import { ColumnType, Experiment } from '../../experiments/webview/contract'
@@ -205,10 +211,26 @@ const getMultiImageInd = (path: string) => {
   return Number(fileName)
 }
 
+const collectImageBoundingBoxes = (
+  boundingBoxes: { label: string; box: BoundingBox }[]
+): { [label: string]: BoundingBox[] } => {
+  const acc: { [label: string]: BoundingBox[] } = {}
+
+  for (const { label, box } of boundingBoxes) {
+    if (!acc[label]) {
+      acc[label] = []
+    }
+
+    acc[label].push(box)
+  }
+
+  return acc
+}
+
 const collectImageData = (
   acc: ComparisonData,
   path: string,
-  plot: ImagePlot
+  plot: ImagePlotOutput
 ) => {
   const isMultiImgPlot = MULTI_IMAGE_PATH_REG.test(path)
   const pathLabel = isMultiImgPlot ? getMultiImagePath(path) : path
@@ -226,12 +248,20 @@ const collectImageData = (
     acc[id][pathLabel] = []
   }
 
-  const imgPlot: ImagePlot = { ...plot }
+  const imgPlot: ImagePlot = {
+    revisions: plot.revisions,
+    type: plot.type,
+    url: plot.url
+  }
 
   if (isMultiImgPlot) {
     imgPlot.ind = getMultiImageInd(path)
   }
 
+  if (plot.boundingBoxes) {
+    imgPlot.boundingBoxes = collectImageBoundingBoxes(plot.boundingBoxes)
+  }
+
   acc[id][pathLabel].push(imgPlot)
 }
 
@@ -312,60 +342,120 @@ export const collectData = (output: PlotsOutput): DataAccumulator => {
   return acc
 }
 
-type ComparisonPlotsAcc = { path: string; revisions: ComparisonRevisionData }[]
+type ComparisonPlotsAcc = {
+  classDetails: ComparisonClassDetails
+  path: string
+  revisions: ComparisonRevisionData
+}[]
 
 type GetComparisonPlotImg = (
   img: ImagePlot,
   id: string,
   path: string
 ) => ComparisonPlotImg
 
+export const boundingBoxColors = [
+  '#ff3838',
+  '#ff9d97',
+  '#ff701f',
+  '#ffb21d',
+  '#cfd231',
+  '#48f90a',
+  '#92cc17',
+  '#3ddb86',
+  '#1a9334',
+  '#00d4bb',
+  '#2c99a8',
+  '#00c2ff',
+  '#344593',
+  '#6473ff',
+  '#0018ec',
+  '#8438ff',
+  '#520085',
+  '#cb38ff',
+  '#ff95c8',
+  '#ff37c7'
+]
+
+const collectSelectedPlotImgClassLabels = (
+  boundingBoxClassLabels: Set<string>,
+  imgs: ImagePlot[] = []
+) => {
+  for (const img of imgs) {
+    if (!img.boundingBoxes) {
+      continue
+    }
+
+    for (const label of Object.keys(img.boundingBoxes)) {
+      boundingBoxClassLabels.add(label)
+    }
+  }
+}
+
 const collectSelectedPathComparisonPlots = ({
   acc,
   comparisonData,
+  comparisonClassesSelected,
   path,
   selectedRevisionIds,
   getComparisonPlotImg
 }: {
   acc: ComparisonPlotsAcc
   comparisonData: ComparisonData
+  comparisonClassesSelected: ComparisonClassesSelected
   path: string
   selectedRevisionIds: string[]
   getComparisonPlotImg: GetComparisonPlotImg
 }) => {
+  const boundingBoxClassLabels = new Set<string>()
   const pathRevisions = {
+    classDetails: {} as ComparisonClassDetails,
     path,
     revisions: {} as ComparisonRevisionData
   }
 
   for (const id of selectedRevisionIds) {
-    const imgs = comparisonData[id]?.[path]
+    const imgs: ImagePlot[] | undefined = comparisonData[id]?.[path]
+
     pathRevisions.revisions[id] = {
       id,
       imgs: imgs
         ? imgs.map(img => getComparisonPlotImg(img, id, path))
         : [{ errors: undefined, loading: false, url: undefined }]
     }
+    collectSelectedPlotImgClassLabels(boundingBoxClassLabels, imgs)
   }
+
+  for (const [ind, label] of [...boundingBoxClassLabels].entries()) {
+    const selectedState = comparisonClassesSelected[path]?.[label]
+    pathRevisions.classDetails[label] = {
+      color: boundingBoxColors[ind % boundingBoxColors.length],
+      selected: selectedState === undefined ? true : selectedState
+    }
+  }
+
   acc.push(pathRevisions)
 }
 
 export const collectSelectedComparisonPlots = ({
   comparisonData,
+  comparisonClassesSelected,
   paths,
   selectedRevisionIds,
   getComparisonPlotImg
 }: {
   comparisonData: ComparisonData
+  comparisonClassesSelected: ComparisonClassesSelected
   paths: string[]
   selectedRevisionIds: string[]
   getComparisonPlotImg: GetComparisonPlotImg
-}) => {
+}): ComparisonPlotsAcc => {
   const acc: ComparisonPlotsAcc = []
 
   for (const path of paths) {
     collectSelectedPathComparisonPlots({
       acc,
+      comparisonClassesSelected,
       comparisonData,
       getComparisonPlotImg,
       path,
@@ -376,6 +466,93 @@ export const collectSelectedComparisonPlots = ({
   return acc
 }
 
+const getSelectedImgComparisonPlotClasses = ({
+  comparisonClassesSelected,
+  img,
+  path
+}: {
+  comparisonClassesSelected: ComparisonClassesSelected
+  img: ImagePlot
+  path: string
+}) => {
+  const imgClasses: ComparisonPlotClass[] = []
+
+  for (const [label, boxes] of Object.entries(img.boundingBoxes || {})) {
+    const selectedState = comparisonClassesSelected[path]?.[label]
+
+    if (selectedState === false) {
+      continue
+    }
+
+    imgClasses.push({
+      boxes,
+      label
+    })
+  }
+
+  return imgClasses
+}
+
+const collectedSelectedPathComparisonPlotClasses = ({
+  acc,
+  comparisonClassesSelected,
+  id,
+  comparisonData,
+  path
+}: {
+  comparisonClassesSelected: ComparisonClassesSelected
+  acc: ComparisonPlotClasses
+  comparisonData: ComparisonData
+  path: string
+  id: string
+}) => {
+  for (const img of comparisonData[id][path]) {
+    const imgClasses = getSelectedImgComparisonPlotClasses({
+      comparisonClassesSelected,
+      img,
+      path
+    })
+
+    if (imgClasses.length === 0) {
+      return
+    }
+
+    if (!acc[id]) {
+      acc[id] = {}
+    }
+
+    acc[id][path] = imgClasses
+  }
+}
+
+export const collectSelectedComparisonPlotClasses = ({
+  comparisonClassesSelected,
+  comparisonData,
+  paths,
+  selectedRevisionIds
+}: {
+  comparisonClassesSelected: ComparisonClassesSelected
+  comparisonData: ComparisonData
+  paths: string[]
+  selectedRevisionIds: string[]
+}) => {
+  const acc: ComparisonPlotClasses = {}
+
+  for (const id of selectedRevisionIds) {
+    for (const path of paths) {
+      collectedSelectedPathComparisonPlotClasses({
+        acc,
+        comparisonClassesSelected,
+        comparisonData,
+        id,
+        path
+      })
+    }
+  }
+
+  return acc
+}
+
 export type TemplateDetailsAccumulator = {
   [path: string]: {
     content: TopLevelSpec