optimize imavid labels rendering in grid, too

voxel51 · Feb 6, 2025 · b506bbb · b506bbb
1 parent a7d8823
commit b506bbb
Showing 1 changed file with 35 additions and 15 deletions.
diff --git a/app/packages/looker/src/lookers/imavid/index.ts b/app/packages/looker/src/lookers/imavid/index.ts
@@ -1,3 +1,5 @@
+import { syncAndGetNewLabels } from "@fiftyone/core/src/components/Sidebar/syncAndGetNewLabels";
+import { gridActivePathsLUT } from "@fiftyone/core/src/components/Sidebar/useDetectNewActiveLabelFields";
 import { BufferManager } from "@fiftyone/utilities";
 import { getImaVidElements } from "../../elements";
 import { IMAVID_SHORTCUTS } from "../../elements/common/actions";
@@ -36,10 +38,32 @@ export class ImaVidLooker extends AbstractLooker<ImaVidState, Sample> {
     }
 
     // subscribe to frame number and update sample when frame number changes
-    this.unsubscribe = this.subscribeToState("currentFrameNumber", () => {
-      this.thisFrameSample?.sample &&
-        this.updateSample(this.thisFrameSample.sample);
-    });
+    this.unsubscribe = this.subscribeToState(
+      "currentFrameNumber",
+      (currentFrameNumber: number) => {
+        const thisFrameId = `${this.uuid}-${currentFrameNumber}`;
+
+        if (
+          gridActivePathsLUT.has(thisFrameId) &&
+          this.thisFrameSample?.sample
+        ) {
+          this.refreshOverlaysToCurrentFrame();
+        } else {
+          const newFieldsIfAny = syncAndGetNewLabels(
+            thisFrameId,
+            gridActivePathsLUT,
+            new Set(this.options.activePaths)
+          );
+
+          if (newFieldsIfAny && currentFrameNumber > 0) {
+            this.refreshSample(newFieldsIfAny, currentFrameNumber);
+          } else {
+            // worst case, only here for fail-safe
+            this.refreshSample();
+          }
+        }
+      }
+    );
   }
 
   get thisFrameSample() {
@@ -207,20 +231,16 @@ export class ImaVidLooker extends AbstractLooker<ImaVidState, Sample> {
   }
 
   refreshOverlaysToCurrentFrame() {
-    const sampleIdFromFramesStore =
-      this.frameStoreController.store.frameIndex.get(this.frameNumber);
+    let { image: _cachedImage, ...thisFrameSample } =
+      this.frameStoreController.store.getSampleAtFrame(
+        this.frameNumber
+      )?.sample;
 
-    let sample: Sample;
-
-    if (sampleIdFromFramesStore) {
-      const { image: _cachedImage, ...sampleWithoutImage } =
-        this.frameStoreController.store.samples.get(sampleIdFromFramesStore);
-      sample = sampleWithoutImage.sample;
-    } else {
-      sample = this.sample;
+    if (!thisFrameSample) {
+      thisFrameSample = this.sample;
     }
 
-    this.loadOverlays(sample);
+    this.loadOverlays(thisFrameSample);
   }
 
   refreshSample(renderLabels: string[] | null = null, frameNumber?: number) {