From 1bc37fefb81e66b4a418808eb0a9b99e4ae2c031 Mon Sep 17 00:00:00 2001
From: Paul Elliott <paul.elliott@kitware.com>
Date: Wed, 23 Oct 2024 10:51:39 -0400
Subject: [PATCH] refactor(paint): consolidate ensure segment group is active
 if exists logic

Rename functions using labelMap to segmentGroup
---
 src/components/SegmentGroupControls.vue |  2 +-
 src/store/tools/paint.ts                | 61 ++++++++++++++-----------
 2 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/src/components/SegmentGroupControls.vue b/src/components/SegmentGroupControls.vue
index 2fcda016..fd8604ef 100644
--- a/src/components/SegmentGroupControls.vue
+++ b/src/components/SegmentGroupControls.vue
@@ -48,7 +48,7 @@ const currentSegmentGroups = computed(() => {
 const paintStore = usePaintToolStore();
 const currentSegmentGroupID = computed({
   get: () => paintStore.activeSegmentGroupID,
-  set: (id) => paintStore.setActiveLabelmap(id),
+  set: (id) => paintStore.setActiveSegmentGroup(id),
 });
 
 // clear selection if we delete the active segment group
diff --git a/src/store/tools/paint.ts b/src/store/tools/paint.ts
index b0350813..6fc74b7d 100644
--- a/src/store/tools/paint.ts
+++ b/src/store/tools/paint.ts
@@ -1,7 +1,7 @@
 import type { Vector2 } from '@kitware/vtk.js/types';
 import { useCurrentImage } from '@/src/composables/useCurrentImage';
 import { Manifest, StateFile } from '@/src/io/state-file/schema';
-import { computed, ref, watch, watchEffect } from 'vue';
+import { computed, ref, watchEffect } from 'vue';
 import { vec3 } from 'gl-matrix';
 import { defineStore } from 'pinia';
 import { Maybe } from '@/src/types';
@@ -48,24 +48,37 @@ export const usePaintToolStore = defineStore('paint', () => {
   /**
    * Sets the active labelmap.
    */
-  function setActiveLabelmap(segmentGroupID: Maybe<string>) {
+  function setActiveSegmentGroup(segmentGroupID: Maybe<string>) {
     activeSegmentGroupID.value = segmentGroupID;
   }
 
+  /**
+   * Gets the first segment group ID for a given image.
+   * @param imageID
+   */
+  function getFirstSegmentGroupID(imageID: Maybe<string>): Maybe<string> {
+    if (!imageID) return null;
+    const segmentGroups = segmentGroupStore.orderByParent[imageID];
+    if (segmentGroups && segmentGroups.length > 0) {
+      return segmentGroups[0];
+    }
+    return null;
+  }
+
   /**
    * Sets the active labelmap from a given image.
    *
    * If a labelmap exists, pick the first one. If no labelmap exists, create one.
    */
-  function setActiveLabelmapFromImage(imageID: Maybe<string>) {
+  function ensureActiveSegmentGroupForImage(imageID: Maybe<string>) {
     if (!imageID) {
-      setActiveLabelmap(null);
+      setActiveSegmentGroup(null);
       return;
     }
 
-    const labelmaps = segmentGroupStore.orderByParent[imageID];
-    if (labelmaps?.length) {
-      activeSegmentGroupID.value = labelmaps[0];
+    const segmentGroupID = getFirstSegmentGroupID(imageID);
+    if (segmentGroupID) {
+      setActiveSegmentGroup(segmentGroupID);
     } else {
       activeSegmentGroupID.value =
         segmentGroupStore.newLabelmapFromImage(imageID);
@@ -157,7 +170,7 @@ export const usePaintToolStore = defineStore('paint', () => {
     if (!imageID) {
       return false;
     }
-    setActiveLabelmapFromImage(imageID);
+    ensureActiveSegmentGroupForImage(imageID);
     this.$paint.setBrushSize(this.brushSize);
 
     isActive.value = true;
@@ -189,28 +202,24 @@ export const usePaintToolStore = defineStore('paint', () => {
     if (paint.activeSegmentGroupID !== null) {
       activeSegmentGroupID.value =
         segmentGroupIDMap[paint.activeSegmentGroupID];
-      setActiveLabelmap(activeSegmentGroupID.value);
+      setActiveSegmentGroup(activeSegmentGroupID.value);
       setActiveSegment.call(this, paint.activeSegment);
     }
   }
 
-  // --- change labelmap if paint is active --- //
+  // Create segment group if paint is active and none exist.
+  // If paint is not active, but there is a segment group for the current image, set it as active.
+  watchEffect(() => {
+    const imageID = currentImageID.value;
+    if (!imageID) return;
 
-  watch(
-    currentImageID,
-    (imageID) => {
-      if (isActive.value) {
-        setActiveLabelmapFromImage(imageID);
+    if (isActive.value) {
+      ensureActiveSegmentGroupForImage(imageID);
+    } else {
+      const segmentGroupID = getFirstSegmentGroupID(imageID);
+      if (segmentGroupID) {
+        setActiveSegmentGroup(segmentGroupID);
       }
-    },
-    { immediate: true }
-  );
-
-  watchEffect(() => {
-    if (!currentImageID.value) return;
-    const segmentGroups = segmentGroupStore.orderByParent[currentImageID.value];
-    if (segmentGroups?.length === 1) {
-      activeSegmentGroupID.value = segmentGroups[0];
     }
   });
 
@@ -229,8 +238,8 @@ export const usePaintToolStore = defineStore('paint', () => {
     deactivateTool,
 
     setMode,
-    setActiveLabelmap,
-    setActiveLabelmapFromImage,
+    setActiveSegmentGroup,
+    ensureActiveSegmentGroupForImage,
     setActiveSegment,
     setBrushSize,
     setSliceAxis,