Add bounding boxes plot frontend components

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

@@ -119,7 +119,8 @@ describe('collectData', () => {
       join('plots', 'acc.png'),
       heatmapPlot,
       join('plots', 'loss.png'),
-      join('plots', 'image')
+      join('plots', 'image'),
+      join('plots', 'bounding_boxes.png')
     ])
 
     const testBranchHeatmap = comparisonData['test-branch'][heatmapPlot]

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

@@ -64,6 +64,14 @@ const plotsDiffFixturePaths: PlotPath[] = [
     revisions: new Set(REVISIONS),
     type: new Set<PathType>([PathType.COMPARISON])
   },
+  {
+    hasChildren: false,
+    label: 'bounding_boxes.png',
+    parentPath: 'plots',
+    path: join('plots', 'bounding_boxes.png'),
+    revisions: new Set(REVISIONS),
+    type: new Set<PathType>([PathType.COMPARISON])
+  },
   {
     hasChildren: false,
     label: 'loss.tsv',
@@ -162,6 +170,7 @@ describe('collectPaths', () => {
       join('plots', 'heatmap.png'),
       join('plots', 'loss.png'),
       join('plots', 'image'),
+      join('plots', 'bounding_boxes.png'),
       join('logs', 'loss.tsv'),
       join('logs', 'acc.tsv'),
       'predictions.json'

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

@@ -81,6 +81,15 @@ describe('PathsModel', () => {
         selected: true,
         type: comparisonType
       },
+      {
+        hasChildren: false,
+        label: 'bounding_boxes.png',
+        parentPath: 'plots',
+        path: join('plots', 'bounding_boxes.png'),
+        revisions: new Set(REVISIONS),
+        selected: true,
+        type: comparisonType
+      },
       {
         hasChildren: false,
         label: 'loss.tsv',
@@ -369,14 +378,16 @@ describe('PathsModel', () => {
       join('plots', 'acc.png'),
       join('plots', 'heatmap.png'),
       join('plots', 'loss.png'),
-      join('plots', 'image')
+      join('plots', 'image'),
+      join('plots', 'bounding_boxes.png')
     ])
 
     const newOrder = [
       join('plots', 'heatmap.png'),
       join('plots', 'acc.png'),
       join('plots', 'loss.png'),
-      join('plots', 'image')
+      join('plots', 'image'),
+      join('plots', 'bounding_boxes.png')
     ]
 
     model.setComparisonPathsOrder(newOrder)
@@ -411,7 +422,7 @@ describe('PathsModel', () => {
         tooltip: undefined
       },
       {
-        descendantStatuses: [2, 2, 2, 2],
+        descendantStatuses: [2, 2, 2, 2, 2],
         hasChildren: true,
         label: 'plots',
         parentPath: undefined,

extension/src/plots/webview/contract.ts

@@ -50,8 +50,13 @@ export type SectionCollapsed = typeof DEFAULT_SECTION_COLLAPSED
 
 export type ComparisonRevisionData = { [revision: string]: ComparisonPlot }
 
+export type ComparisonBoundingBoxLabels = {
+  [label: string]: { selected: boolean; color: string }
+}
+
 export type ComparisonPlots = {
   path: string
+  boundingBoxLabels: ComparisonBoundingBoxLabels
   revisions: ComparisonRevisionData
 }[]
 
@@ -151,11 +156,22 @@ export interface TemplatePlotsData {
   smoothPlotValues: SmoothPlotValues
 }
 
+export type ComparisonPlotBoundingBoxes = {
+  label: string
+  boxes: {
+    h: number
+    w: number
+    x: number
+    y: number
+  }[]
+}[]
+
 export type ComparisonPlotImg = {
   url: string | undefined
   errors: string[] | undefined
   loading: boolean
   ind?: number
+  boundingBoxes?: ComparisonPlotBoundingBoxes
 }
 
 export type ComparisonPlot = {

extension/src/plots/webview/messages.ts

@@ -406,7 +406,11 @@ export class WebviewMessages {
       height: this.plots.getHeight(PlotsSection.COMPARISON_TABLE),
       multiPlotValues: this.plots.getComparisonMultiPlotValues(),
       plots: comparison.map(({ path, revisions }) => {
-        return { path, revisions: this.getRevisionsWithCorrectUrls(revisions) }
+        return {
+          boundingBoxLabels: {},
+          path,
+          revisions: this.getRevisionsWithCorrectUrls(revisions)
+        }
       }),
       revisions: this.plots.getComparisonRevisions(),
       width: this.plots.getNbItemsPerRowOrWidth(PlotsSection.COMPARISON_TABLE)

extension/src/test/fixtures/plotsDiff/index.ts

@@ -18,7 +18,8 @@ import {
   DEFAULT_PLOT_HEIGHT,
   DEFAULT_NB_ITEMS_PER_ROW,
   DEFAULT_PLOT_WIDTH,
-  ComparisonPlotImg
+  ComparisonPlotImg,
+  ComparisonBoundingBoxLabels
 } from '../../../plots/webview/contract'
 import { join } from '../../util/path'
 import { copyOriginalColors } from '../../../experiments/model/status/colors'
@@ -626,7 +627,34 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
     'exp-e7a67',
     'test-branch',
     'exp-83425'
-  ])
+  ]),
+  [join('plots', 'bounding_boxes.png')]: [
+    {
+      type: PlotsType.IMAGE,
+      revisions: [EXPERIMENT_WORKSPACE_ID],
+      url: joinFunc(baseUrl, 'bounding_boxes.png')
+    },
+    {
+      type: PlotsType.IMAGE,
+      revisions: ['main'],
+      url: joinFunc(baseUrl, 'bounding_boxes.png')
+    },
+    {
+      type: PlotsType.IMAGE,
+      revisions: ['exp-e7a67'],
+      url: joinFunc(baseUrl, 'bounding_boxes.png')
+    },
+    {
+      type: PlotsType.IMAGE,
+      revisions: ['test-branch'],
+      url: joinFunc(baseUrl, 'bounding_boxes.png')
+    },
+    {
+      type: PlotsType.IMAGE,
+      revisions: ['exp-83425'],
+      url: joinFunc(baseUrl, 'bounding_boxes.png')
+    }
+  ]
 })
 
 export const getOutput = (baseUrl: string): PlotsOutput => ({
@@ -913,17 +941,23 @@ export const getComparisonWebviewMessage = (
   joinFunc: (...args: string[]) => string = join
 ): PlotsComparisonData => {
   const plotAcc: {
-    [path: string]: { path: string; revisions: ComparisonRevisionData }
+    [path: string]: {
+      path: string
+      revisions: ComparisonRevisionData
+      boundingBoxLabels: ComparisonBoundingBoxLabels
+    }
   } = {}
 
   for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
     const isMulti = path.includes('image')
+    const isBoundingBox = path.includes('bounding_boxes.png')
     const pathLabel = isMulti ? join('plots', 'image') : path
 
     if (!plotAcc[pathLabel]) {
       plotAcc[pathLabel] = {
         path: pathLabel,
-        revisions: {}
+        revisions: {},
+        boundingBoxLabels: {}
       }
     }
 

extension/src/test/suite/plots/index.test.ts

@@ -402,7 +402,8 @@ suite('Plots Test Suite', () => {
         join('plots', 'acc.png'),
         join('plots', 'heatmap.png'),
         join('plots', 'loss.png'),
-        join('plots', 'image')
+        join('plots', 'image'),
+        join('plots', 'bounding_boxes.png')
       ]
 
       messageSpy.resetHistory()

webview/src/plots/components/App.test.tsx

@@ -228,6 +228,7 @@ describe('App', () => {
         multiPlotValues: {},
         plots: [
           {
+            boundingBoxLabels: {},
             path: 'training/plots/images/misclassified.jpg',
             revisions: {
               ad2b5ec: {
@@ -283,6 +284,7 @@ describe('App', () => {
         multiPlotValues: {},
         plots: [
           {
+            boundingBoxLabels: {},
             path: 'training/plots/images/image',
             revisions: {
               ad2b5ec: {

webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx

@@ -30,6 +30,7 @@ import { plotsReducers, plotsStore } from '../../store'
 import { webviewInitialState } from '../webviewSlice'
 import { getThemeValue, hexToRGB, ThemeProperty } from '../../../util/styles'
 import * as EventCurrentTargetDistances from '../../../shared/components/dragDrop/currentTarget'
+import { addBoundingBoxes } from '../../../test/boundingBoxesFixture'
 
 const getHeaders = (): HTMLElement[] => screen.getAllByRole('columnheader')
 
@@ -297,8 +298,8 @@ describe('ComparisonTable', () => {
 
     renderTable({
       ...comparisonTableFixture,
-      plots: comparisonTableFixture.plots.map(({ path, revisions }) => ({
-        path,
+      plots: comparisonTableFixture.plots.map(({ revisions, ...rest }) => ({
+        ...rest,
         revisions: {
           ...revisions,
           [revisionWithNoData]: {
@@ -334,8 +335,8 @@ describe('ComparisonTable', () => {
 
     renderTable({
       ...comparisonTableFixture,
-      plots: comparisonTableFixture.plots.map(({ path, revisions }) => ({
-        path,
+      plots: comparisonTableFixture.plots.map(({ revisions, ...rest }) => ({
+        ...rest,
         revisions: {
           ...revisions,
           [revisionWithNoData]: {
@@ -743,4 +744,44 @@ describe('ComparisonTable', () => {
       })
     })
   })
+
+  describe('Plots With Bounding Boxes', () => {
+    const plotsWithBoundingBoxes = addBoundingBoxes(comparisonTableFixture)
+
+    it('should show toggable labels in the plot row', () => {
+      renderTable(plotsWithBoundingBoxes)
+
+      const rowHeaders = screen.getAllByTestId('row-header')
+
+      const boundingBoxPlotHeader = rowHeaders[4]
+
+      expect(
+        within(boundingBoxPlotHeader).getByText('Labels')
+      ).toBeInTheDocument()
+
+      const checkedLabel = within(boundingBoxPlotHeader).getByLabelText(
+        'traffic light'
+      )
+      expect(checkedLabel).toBeInTheDocument()
+      expect(checkedLabel).toHaveAttribute('checked')
+      const uncheckedLabel = within(boundingBoxPlotHeader).getByLabelText(
+        'sign'
+      )
+      expect(uncheckedLabel).toBeInTheDocument()
+      expect(uncheckedLabel).not.toHaveAttribute('checked')
+    })
+
+    it('should show svgs with bounding boxes instead of images', () => {
+      renderTable(plotsWithBoundingBoxes)
+
+      const boundingBoxPlotImage = screen.getByLabelText(
+        /bounding_boxes.png \(workspace\)/
+      )
+      expect(boundingBoxPlotImage).toHaveAttribute('viewBox')
+      expect(
+        within(boundingBoxPlotImage).getByText('traffic light')
+      ).toBeInTheDocument()
+      expect(within(boundingBoxPlotImage).getAllByText('car')).toHaveLength(2)
+    })
+  })
 })

webview/src/plots/components/comparisonTable/ComparisonTableRow.test.tsx

@@ -20,6 +20,7 @@ jest.mock('../../../shared/api')
 
 describe('ComparisonTableRow', () => {
   const basicProps: ComparisonTableRowProps = {
+    boundingBoxLabels: {},
     nbColumns: 3,
     onLayoutChange: jest.fn(),
     order: ['path/to/the-file/image.png'],