Fix "Image By Step" step number inaccuracy (#4516)

Author: julieg18

extension/src/plots/model/collect.ts

@@ -134,11 +134,9 @@ export type RevisionData = {
   [label: string]: RevisionPathData
 }
 
-type ComparisonDataImgPlot = ImagePlot & { ind?: number }
-
 export type ComparisonData = {
   [label: string]: {
-    [path: string]: ComparisonDataImgPlot[]
+    [path: string]: ImagePlot[]
   }
 }
 
@@ -171,7 +169,7 @@ const collectImageData = (
     acc[id][pathLabel] = []
   }
 
-  const imgPlot: ComparisonDataImgPlot = { ...plot }
+  const imgPlot: ImagePlot = { ...plot }
 
   if (isMultiImgPlot) {
     imgPlot.ind = getMultiImageInd(path)

extension/src/plots/model/index.ts

@@ -34,7 +34,8 @@ import {
   PlotHeight,
   SmoothPlotValues,
   ImagePlot,
-  ComparisonMultiPlotValues
+  ComparisonMultiPlotValues,
+  ComparisonPlotImg
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -462,11 +463,17 @@ export class PlotsModel extends ModelWithPersistence {
         const url = collectImageUrl(image, fetched)
         const loading = !fetched && !url
 
-        return {
+        const plot: ComparisonPlotImg = {
           errors,
           loading,
           url
         }
+
+        if (typeof image.ind === 'number') {
+          plot.ind = image.ind
+        }
+
+        return plot
       },
       paths,
       selectedRevisionIds

extension/src/plots/webview/contract.ts

@@ -124,6 +124,7 @@ export type ImagePlot = {
   revisions?: string[]
   type: PlotsType
   url: string
+  ind?: number
 }
 
 export const isImagePlot = (plot: Plot): plot is ImagePlot =>
@@ -158,6 +159,7 @@ export type ComparisonPlotImg = {
   url: string | undefined
   errors: string[] | undefined
   loading: boolean
+  ind?: number
 }
 
 export type ComparisonPlot = {

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

@@ -14,7 +14,8 @@ import {
   PlotsComparisonData,
   DEFAULT_PLOT_HEIGHT,
   DEFAULT_NB_ITEMS_PER_ROW,
-  DEFAULT_PLOT_WIDTH
+  DEFAULT_PLOT_WIDTH,
+  ComparisonPlotImg
 } from '../../../plots/webview/contract'
 import { join } from '../../util/path'
 import { copyOriginalColors } from '../../../experiments/model/status/colors'
@@ -375,10 +376,14 @@ const getMultiImageData = (
       url: string
     }[]
   } = {}
+
   for (let i = 0; i < 15; i++) {
     const key = joinFunc('plots', 'image', `${i}.jpg`)
     const values = []
     for (const revision of revisions) {
+      if (revision === 'exp-83425' && i % 2 === 0) {
+        continue
+      }
       values.push({
         type: PlotsType.IMAGE,
         revisions: [revision],
@@ -770,6 +775,11 @@ export const getManyTemplatePlotsWebviewMessage = (
 
 export const MOCK_IMAGE_MTIME = 946684800000
 
+const getIndFromComparisonMultiImgPath = (path: string) => {
+  const pathIndMatches = path.match(/(\d+)\.jpg$/)
+  return Number((pathIndMatches as string[])[1])
+}
+
 export const getComparisonWebviewMessage = (
   baseUrl: string,
   joinFunc: (...args: string[]) => string = join
@@ -779,7 +789,8 @@ export const getComparisonWebviewMessage = (
   } = {}
 
   for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
-    const pathLabel = path.includes('image') ? join('plots', 'image') : path
+    const isMulti = path.includes('image')
+    const pathLabel = isMulti ? join('plots', 'image') : path
 
     if (!plotAcc[pathLabel]) {
       plotAcc[pathLabel] = {
@@ -801,11 +812,17 @@ export const getComparisonWebviewMessage = (
         }
       }
 
-      plotAcc[pathLabel].revisions[id].imgs.push({
+      const img: ComparisonPlotImg = {
         url: `${url}?${MOCK_IMAGE_MTIME}`,
         errors: undefined,
         loading: false
-      })
+      }
+
+      if (isMulti) {
+        img.ind = getIndFromComparisonMultiImgPath(path)
+      }
+
+      plotAcc[pathLabel].revisions[id].imgs.push(img)
     }
   }
 

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

@@ -1880,6 +1880,37 @@ describe('App', () => {
         within(multiImgPlots[1]).queryByRole('img')
       ).not.toBeInTheDocument()
     })
+
+    it('should handle a plot with steps that do not increment by one', () => {
+      renderAppWithOptionalData({
+        comparison: comparisonTableFixture
+      })
+
+      const unusualRev = 'exp-83425'
+      const imgs = comparisonTableFixture.plots[3].revisions[unusualRev].imgs
+      const multiImgPlot = screen.getAllByTestId('multi-image-cell')[4]
+
+      const slider = within(multiImgPlot).getByRole('slider')
+      const imgEl = within(multiImgPlot).getByRole('img')
+
+      expect(slider).toHaveAttribute('max', '6')
+
+      expect(imgEl).toHaveAttribute('src', imgs[0].url)
+      expect(imgEl).toHaveAttribute(
+        'alt',
+        `1 of ${join('plots', 'image')} (exp-83425)`
+      )
+      expect(within(multiImgPlot).getByText('1')).toBeInTheDocument()
+
+      fireEvent.change(slider, { target: { value: 3 } })
+
+      expect(imgEl).toHaveAttribute('src', imgs[3].url)
+      expect(imgEl).toHaveAttribute(
+        'alt',
+        `7 of ${join('plots', 'image')} (exp-83425)`
+      )
+      expect(within(multiImgPlot).getByText('7')).toBeInTheDocument()
+    })
   })
 
   describe('Virtualization', () => {

webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx

@@ -1,6 +1,9 @@
 import React, { useEffect, useCallback, useRef, useState } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
-import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
+import {
+  ComparisonPlot,
+  ComparisonPlotImg
+} from 'dvc/src/plots/webview/contract'
 import { ComparisonTableCell } from './ComparisonTableCell'
 import styles from '../styles.module.scss'
 import { changeDisabledDragIds } from '../comparisonTableSlice'
@@ -18,6 +21,12 @@ export const ComparisonTableMultiCell: React.FC<{
   const dispatch = useDispatch()
   const maxStep = plot.imgs.length - 1
   const changeDebounceTimer = useRef(0)
+  const selectedImg: ComparisonPlotImg = plot.imgs[currentStep] || {
+    errors: undefined,
+    ind: currentStep,
+    loading: false,
+    url: undefined
+  }
 
   const addDisabled = useCallback(() => {
     dispatch(changeDisabledDragIds([path]))
@@ -46,15 +55,9 @@ export const ComparisonTableMultiCell: React.FC<{
         path={path}
         plot={{
           id: plot.id,
-          imgs: [
-            plot.imgs[currentStep] || {
-              errors: undefined,
-              loading: false,
-              url: undefined
-            }
-          ]
+          imgs: [selectedImg]
         }}
-        imgAlt={`${currentStep} of ${path} (${plot.id})`}
+        imgAlt={`${selectedImg.ind} of ${path} (${plot.id})`}
       />
       <div
         className={styles.multiImageSlider}
@@ -76,7 +79,7 @@ export const ComparisonTableMultiCell: React.FC<{
             setCurrentStep(Number(event.target.value))
           }}
         />
-        <p>{currentStep}</p>
+        <p>{selectedImg.ind}</p>
       </div>
     </div>
   )

webview/src/stories/ComparisonTable.stories.tsx

@@ -88,11 +88,15 @@ const removeImages = (
       filteredRevisionData[id] = {
         id,
         imgs: isMulti
-          ? (Array.from({ length: revisionsData[id].imgs.length }).fill({
-              errors: undefined,
-              loading: false,
-              url: undefined
-            }) as ComparisonPlotImg[])
+          ? Array.from({ length: revisionsData[id].imgs.length }).map(
+              (_, ind) =>
+                ({
+                  errors: undefined,
+                  ind,
+                  loading: false,
+                  url: undefined
+                }) as ComparisonPlotImg
+            )
           : [
               {
                 errors: