From 2c6ce21c8f79f8b3e6dc47b5aa6f717b6b5cbd2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vin=C3=ADcius=20Alves=20de=20Faria=20Resende?=
 <viniciusfariaresende@gmail.com>
Date: Tue, 7 Oct 2025 14:23:26 -0300
Subject: [PATCH] fix(Segmentation): [Bug #5420] Segmentation color resets
 after using toggleOneUp (#5465)

---
 .../SegmentationService.test.ts               | 59 +++++++++++++++++--
 .../SegmentationService.ts                    |  6 ++
 2 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/extensions/cornerstone/src/services/SegmentationService/SegmentationService.test.ts b/extensions/cornerstone/src/services/SegmentationService/SegmentationService.test.ts
index 0b3e6de62..eed05265f 100644
--- a/extensions/cornerstone/src/services/SegmentationService/SegmentationService.test.ts
+++ b/extensions/cornerstone/src/services/SegmentationService/SegmentationService.test.ts
@@ -2031,6 +2031,9 @@ describe('SegmentationService', () => {
         .spyOn(cstSegmentation.segmentLocking, 'setSegmentIndexLocked')
         .mockReturnValue(undefined);
       jest.spyOn(cstSegmentation.config.color, 'setSegmentIndexColor').mockReturnValue(undefined);
+      jest
+        .spyOn(cstSegmentation.state, 'getSegmentationRepresentations')
+        .mockReturnValue([{ colorLUTIndex: 1 }]);
 
       service.addSegment(segmentationId, config);
 
@@ -2192,16 +2195,25 @@ describe('SegmentationService', () => {
   });
 
   describe('setSegmentColor', () => {
-    it('should set the color of the segment', () => {
-      const viewportId = 'viewportId';
-      const segmentationId = 'segmentationId';
-      const segmentIndex = 1;
-      const color = [255, 0, 0, 255] as csTypes.Color;
+    const viewportId = 'viewportId';
+    const segmentationId = 'segmentationId';
+    const segmentIndex = 1;
+    const color = [255, 0, 0, 255] as csTypes.Color;
 
+    it('should set the color of the segment', () => {
+      jest
+        .spyOn(cstSegmentation.state, 'getSegmentationRepresentations')
+        .mockReturnValue([{ colorLUTIndex: 1 }]);
       jest.spyOn(cstSegmentation.config.color, 'setSegmentIndexColor').mockReturnValue(undefined);
 
       service.setSegmentColor(viewportId, segmentationId, segmentIndex, color);
 
+      expect(cstSegmentation.state.getSegmentationRepresentations).toHaveBeenCalledTimes(1);
+      expect(cstSegmentation.state.getSegmentationRepresentations).toHaveBeenCalledWith(
+        viewportId,
+        { segmentationId }
+      );
+
       expect(cstSegmentation.config.color.setSegmentIndexColor).toHaveBeenCalledTimes(1);
       expect(cstSegmentation.config.color.setSegmentIndexColor).toHaveBeenCalledWith(
         viewportId,
@@ -2210,6 +2222,43 @@ describe('SegmentationService', () => {
         color
       );
     });
+
+    it('should set the color of the segment with the colorLUTIndex', async () => {
+      jest
+        .spyOn(cstSegmentation.state, 'getSegmentationRepresentations')
+        .mockReturnValue([{ colorLUTIndex: 1 }]);
+      jest.spyOn(cstSegmentation.config.color, 'setSegmentIndexColor').mockReturnValue(undefined);
+
+      service.setSegmentColor(viewportId, segmentationId, segmentIndex, color);
+
+      jest
+        .spyOn(cstSegmentation.state, 'getSegmentation')
+        .mockReturnValue(mockCornerstoneSegmentation as cstTypes.Segmentation);
+      jest
+        .spyOn(serviceManagerMock.services.cornerstoneViewportService, 'getCornerstoneViewport')
+        // only needed interfaces for the addSegmentationRepresentation call
+        .mockReturnValue(mockCornerstoneStackViewport as unknown as csTypes.IStackViewport);
+      jest
+        .spyOn(cstSegmentation.state, 'updateLabelmapSegmentationImageReferences')
+        .mockReturnValue('labelmapImageId');
+      jest.spyOn(cstSegmentation, 'addSegmentationRepresentations').mockReturnValueOnce(undefined);
+
+      await service.addSegmentationRepresentation(viewportId, {
+        segmentationId: segmentationId,
+        type: csToolsEnums.SegmentationRepresentations.Labelmap,
+        config: { active: true },
+        suppressEvents: true,
+      });
+
+      expect(cstSegmentation.addSegmentationRepresentations).toHaveBeenCalledTimes(1);
+      expect(cstSegmentation.addSegmentationRepresentations).toHaveBeenCalledWith(viewportId, [
+        {
+          type: csToolsEnums.SegmentationRepresentations.Labelmap,
+          segmentationId: segmentationId,
+          config: { colorLUTOrIndex: 1, active: true },
+        },
+      ]);
+    });
   });
 
   describe('getSegmentColor', () => {
diff --git a/extensions/cornerstone/src/services/SegmentationService/SegmentationService.ts b/extensions/cornerstone/src/services/SegmentationService/SegmentationService.ts
index 6368ef291..6946fd897 100644
--- a/extensions/cornerstone/src/services/SegmentationService/SegmentationService.ts
+++ b/extensions/cornerstone/src/services/SegmentationService/SegmentationService.ts
@@ -983,6 +983,12 @@ class SegmentationService extends PubSubService {
     segmentIndex: number,
     color: csTypes.Color
   ): void {
+    const segmentationRepresentations = this.getSegmentationRepresentations(viewportId, {
+      segmentationId,
+    });
+    const { colorLUTIndex } = segmentationRepresentations[0];
+    this._segmentationIdToColorLUTIndexMap.set(segmentationId, colorLUTIndex);
+
     cstSegmentation.config.color.setSegmentIndexColor(
       viewportId,
       segmentationId,