Add "Image by Step" plots

extension/src/cli/dvc/constants.ts

@@ -18,6 +18,8 @@ export const EXP_RWLOCK_FILE = join(TEMP_EXP_DIR, 'rwlock.lock')
 export const DEFAULT_NUM_OF_COMMITS_TO_SHOW = 3
 export const NUM_OF_COMMITS_TO_INCREASE = 2
 
+export const MULTI_IMAGE_PATH_REG = /[^/]+[/\\]\d+\.[a-z]+$/i
+
 export enum Command {
   ADD = 'add',
   CHECKOUT = 'checkout',

extension/src/cli/dvc/index.test.ts

@@ -1,7 +1,8 @@
+import { join } from 'path'
 import { EventEmitter } from 'vscode'
 import { Disposable, Disposer } from '@hediet/std/disposable'
 import { DvcCli } from '.'
-import { Command } from './constants'
+import { Command, MULTI_IMAGE_PATH_REG } from './constants'
 import { CliResult, CliStarted, typeCheckCommands } from '..'
 import { getProcessEnv } from '../../env'
 import { createProcess } from '../../process/execution'
@@ -52,6 +53,53 @@ describe('typeCheckCommands', () => {
   })
 })
 
+describe('Comparison Multi Image Regex', () => {
+  it('should match a nested image group directory', () => {
+    expect(
+      MULTI_IMAGE_PATH_REG.test(
+        join(
+          'extremely',
+          'super',
+          'super',
+          'super',
+          'nested',
+          'image',
+          '768.svg'
+        )
+      )
+    ).toBe(true)
+  })
+
+  it('should match directories with spaces or special characters', () => {
+    expect(MULTI_IMAGE_PATH_REG.test(join('mis classified', '5.png'))).toBe(
+      true
+    )
+
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified#^', '5.png'))).toBe(
+      true
+    )
+  })
+
+  it('should match different types of images', () => {
+    const imageFormats = ['svg', 'png', 'jpg', 'jpeg']
+    for (const format of imageFormats) {
+      expect(
+        MULTI_IMAGE_PATH_REG.test(join('misclassified', `5.${format}`))
+      ).toBe(true)
+    }
+  })
+
+  it('should not match files that include none digits or do not have a file extension', () => {
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', 'five.png'))).toBe(
+      false
+    )
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', '5 4.png'))).toBe(
+      false
+    )
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', '5'))).toBe(false)
+  })
+})
+
 describe('executeDvcProcess', () => {
   it('should pass the correct details to the underlying process given no path to the cli or python binary path', async () => {
     const existingPath = joinEnvPath(

extension/src/fileSystem/util.ts

@@ -1,4 +1,4 @@
-import { sep } from 'path'
+import { sep, parse } from 'path'
 
 export const getPathArray = (path: string): string[] => path.split(sep)
 
@@ -18,3 +18,5 @@ export const getParent = (pathArray: string[], idx: number) => {
 
 export const removeTrailingSlash = (path: string): string =>
   path.endsWith(sep) ? path.slice(0, -1) : path
+
+export const getFileNameWithoutExt = (path: string) => parse(path).name

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

@@ -84,7 +84,8 @@ describe('collectData', () => {
     expect(Object.keys(comparisonData.main)).toStrictEqual([
       join('plots', 'acc.png'),
       heatmapPlot,
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      join('plots', 'image')
     ])
 
     const testBranchHeatmap = comparisonData['test-branch'][heatmapPlot]

extension/src/plots/model/collect.ts

@@ -36,6 +36,12 @@ import {
 import { StrokeDashEncoding } from '../multiSource/constants'
 import { exists } from '../../fileSystem'
 import { hasKey } from '../../util/object'
+import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
+import {
+  getFileNameWithoutExt,
+  getParent,
+  getPathArray
+} from '../../fileSystem/util'
 
 export const getCustomPlotId = (metric: string, param: string) =>
   `custom-${metric}-${param}`
@@ -128,19 +134,31 @@ export type RevisionData = {
   [label: string]: RevisionPathData
 }
 
+type ComparisonDataImgPlot = ImagePlot & { ind?: number }
+
 export type ComparisonData = {
   [label: string]: {
-    [path: string]: ImagePlot[]
+    [path: string]: ComparisonDataImgPlot[]
   }
 }
 
+const getMultiImagePath = (path: string) =>
+  getParent(getPathArray(path), 0) as string
+
+const getMultiImageInd = (path: string) => {
+  const fileName = getFileNameWithoutExt(path)
+  return Number(fileName)
+}
+
 const collectImageData = (
   acc: ComparisonData,
   path: string,
   plot: ImagePlot
 ) => {
-  const pathLabel = path
+  const isMultiImgPlot = MULTI_IMAGE_PATH_REG.test(path)
+  const pathLabel = isMultiImgPlot ? getMultiImagePath(path) : path
   const id = plot.revisions?.[0]
+
   if (!id) {
     return
   }
@@ -153,7 +171,13 @@ const collectImageData = (
     acc[id][pathLabel] = []
   }
 
-  acc[id][pathLabel].push(plot)
+  const imgPlot: ComparisonDataImgPlot = { ...plot }
+
+  if (isMultiImgPlot) {
+    imgPlot.ind = getMultiImageInd(path)
+  }
+
+  acc[id][pathLabel].push(imgPlot)
 }
 
 const collectDatapoints = (
@@ -209,6 +233,16 @@ const collectPathData = (acc: DataAccumulator, path: string, plots: Plot[]) => {
   }
 }
 
+const sortComparisonImgPaths = (acc: DataAccumulator) => {
+  for (const [label, paths] of Object.entries(acc.comparisonData)) {
+    for (const path of Object.keys(paths)) {
+      acc.comparisonData[label][path].sort(
+        (img1, img2) => (img1.ind || 0) - (img2.ind || 0)
+      )
+    }
+  }
+}
+
 export const collectData = (output: PlotsOutput): DataAccumulator => {
   const { data } = output
   const acc = {
@@ -220,6 +254,8 @@ export const collectData = (output: PlotsOutput): DataAccumulator => {
     collectPathData(acc, path, plots)
   }
 
+  sortComparisonImgPaths(acc)
+
   return acc
 }
 
@@ -248,7 +284,7 @@ const collectSelectedPathComparisonPlots = ({
   }
 
   for (const id of selectedRevisionIds) {
-    const imgs = comparisonData[id][path]
+    const imgs = comparisonData[id]?.[path]
     pathRevisions.revisions[id] = {
       id,
       imgs: imgs

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

@@ -45,6 +45,13 @@ describe('collectPaths', () => {
         revisions: new Set(REVISIONS),
         type: new Set(['comparison'])
       },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'image'),
+        revisions: new Set(REVISIONS),
+        type: new Set(['comparison'])
+      },
       {
         hasChildren: false,
         parentPath: 'logs',

extension/src/plots/paths/collect.ts

@@ -23,6 +23,7 @@ import {
 } from '../multiSource/constants'
 import { MultiSourceEncoding } from '../multiSource/collect'
 import { truncate } from '../../util/string'
+import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
 
 export enum PathType {
   COMPARISON = 'comparison',
@@ -58,8 +59,13 @@ const collectType = (plots: Plot[]) => {
 const getType = (
   data: PlotsData,
   hasChildren: boolean,
-  path: string
+  path: string,
+  isMultiImgPlot?: boolean
 ): Set<PathType> | undefined => {
+  if (isMultiImgPlot) {
+    return new Set<PathType>([PathType.COMPARISON])
+  }
+
   if (hasChildren) {
     return
   }
@@ -123,7 +129,8 @@ const collectOrderedPath = (
   data: PlotsData,
   revisions: Set<string>,
   pathArray: string[],
-  idx: number
+  idx: number,
+  isMultiImgDir: boolean
 ): PlotPath[] => {
   const path = getPath(pathArray, idx)
 
@@ -147,7 +154,10 @@ const collectOrderedPath = (
     revisions
   }
 
-  const type = getType(data, hasChildren, path)
+  const isPathLeaf = idx === pathArray.length
+  const isMultiImgPlot = isMultiImgDir && isPathLeaf
+
+  const type = getType(data, hasChildren, path, isMultiImgPlot)
   if (type) {
     plotPath.type = type
   }
@@ -167,9 +177,23 @@ const addRevisionsToPath = (
   }
 
   const pathArray = getPathArray(path)
+  const isMultiImg =
+    MULTI_IMAGE_PATH_REG.test(path) &&
+    !!getType(data, false, path)?.has(PathType.COMPARISON)
+
+  if (isMultiImg) {
+    pathArray.pop()
+  }
 
   for (let reverseIdx = pathArray.length; reverseIdx > 0; reverseIdx--) {
-    acc = collectOrderedPath(acc, data, revisions, pathArray, reverseIdx)
+    acc = collectOrderedPath(
+      acc,
+      data,
+      revisions,
+      pathArray,
+      reverseIdx,
+      isMultiImg
+    )
   }
   return acc
 }

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

@@ -59,6 +59,14 @@ describe('PathsModel', () => {
         selected: true,
         type: comparisonType
       },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'image'),
+        revisions: new Set(REVISIONS),
+        selected: true,
+        type: comparisonType
+      },
       {
         hasChildren: false,
         parentPath: 'logs',
@@ -340,13 +348,15 @@ describe('PathsModel', () => {
     expect(model.getComparisonPaths()).toStrictEqual([
       join('plots', 'acc.png'),
       join('plots', 'heatmap.png'),
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      join('plots', 'image')
     ])
 
     const newOrder = [
       join('plots', 'heatmap.png'),
       join('plots', 'acc.png'),
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      join('plots', 'image')
     ]
 
     model.setComparisonPathsOrder(newOrder)
@@ -380,7 +390,7 @@ describe('PathsModel', () => {
         tooltip: undefined
       },
       {
-        descendantStatuses: [2, 2, 2],
+        descendantStatuses: [2, 2, 2, 2],
         hasChildren: true,
         parentPath: undefined,
         path: 'plots',

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

@@ -377,7 +377,7 @@ const getMultiImageData = (
     }[]
   } = {}
   for (let i = 0; i < 15; i++) {
-    const key = join('plots', 'image', `${i}.jpg`)
+    const key = joinFunc('plots', 'image', `${i}.jpg`)
     const values = []
     for (const revision of revisions) {
       values.push({
@@ -472,11 +472,7 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
       revisions: ['exp-83425'],
       url: joinFunc(baseUrl, '1ba7bcd_plots_loss.png')
     }
-  ]
-})
-
-const getImageDataWithMultiImgs = (baseUrl: string, joinFunc = join) => ({
-  ...getImageData(baseUrl, joinFunc),
+  ],
   ...getMultiImageData(baseUrl, joinFunc, [
     EXPERIMENT_WORKSPACE_ID,
     'main',
@@ -797,21 +793,14 @@ export const MOCK_IMAGE_MTIME = 946684800000
 
 export const getComparisonWebviewMessage = (
   baseUrl: string,
-  joinFunc: (...args: string[]) => string = join,
-  addMulti?: boolean
+  joinFunc: (...args: string[]) => string = join
 ): PlotsComparisonData => {
   const plotAcc: {
     [path: string]: { path: string; revisions: ComparisonRevisionData }
   } = {}
 
-  for (const [path, plots] of Object.entries(
-    addMulti
-      ? getImageDataWithMultiImgs(baseUrl, joinFunc)
-      : getImageData(baseUrl, joinFunc)
-  )) {
-    const multiImagePath = joinFunc('plots', 'image')
-    const isMulti = path.includes(multiImagePath)
-    const pathLabel = path
+  for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
+    const pathLabel = path.includes('image') ? join('plots', 'image') : path
 
     if (!plotAcc[pathLabel]) {
       plotAcc[pathLabel] = {

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

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