diff --git a/src/datasource/graphene/frontend.ts b/src/datasource/graphene/frontend.ts index 0a7e7a1e97..1bf77666ad 100644 --- a/src/datasource/graphene/frontend.ts +++ b/src/datasource/graphene/frontend.ts @@ -57,6 +57,11 @@ import { PYCG_APP_VERSION, getHttpSource, } from "#src/datasource/graphene/base.js"; +import { + applySplitPreviewToTemporaryState, + MulticutSplitPreviewState, + parseGrapheneSplitPreviewResponse, +} from "#src/datasource/graphene/split_preview.js"; import type { DataSource, DataSourceLookupResult, @@ -1153,6 +1158,7 @@ class MulticutState extends RefCounted implements Trackable { sinks = new WatchableSet(); sources = new WatchableSet(); + preview = new MulticutSplitPreviewState(); constructor( public focusSegment = new TrackableValue( @@ -1169,13 +1175,28 @@ class MulticutState extends RefCounted implements Trackable { } }; - this.registerDisposer(focusSegment.changed.add(this.changed.dispatch)); + this.registerDisposer( + focusSegment.changed.add(() => { + this.preview.invalidate(); + this.changed.dispatch(); + }), + ); this.registerDisposer(this.sinks.changed.add(maybeResetFocusSegemnt)); this.registerDisposer(this.sources.changed.add(maybeResetFocusSegemnt)); this.registerDisposer(this.blueGroup.changed.add(this.changed.dispatch)); - this.registerDisposer(this.sinks.changed.add(this.changed.dispatch)); - this.registerDisposer(this.sources.changed.add(this.changed.dispatch)); + this.registerDisposer( + this.sinks.changed.add(() => { + this.preview.invalidate(); + this.changed.dispatch(); + }), + ); + this.registerDisposer( + this.sources.changed.add(() => { + this.preview.invalidate(); + this.changed.dispatch(); + }), + ); } replaceSegments(oldValues: Uint64Set, newValues: Uint64Set) { @@ -1183,6 +1204,13 @@ class MulticutState extends RefCounted implements Trackable { const { focusSegment: { value: focusSegment }, } = this; + const rootsChanged = + (focusSegment !== undefined && oldValues.has(focusSegment)) || + [...this.sinks].some((sink) => oldValues.has(sink.rootId)) || + [...this.sources].some((source) => oldValues.has(source.rootId)); + if (rootsChanged) { + this.preview.invalidate(); + } if (focusSegment && oldValues.has(focusSegment)) { if (newValue) { this.focusSegment.value = newValue; @@ -1200,12 +1228,30 @@ class MulticutState extends RefCounted implements Trackable { } reset() { + this.preview.invalidate(); this.focusSegment.reset(); this.blueGroup.value = false; this.sinks.clear(); this.sources.clear(); } + setPreviewPending(value: boolean) { + if (this.preview.setPending(value)) { + this.changed.dispatch(); + } + } + + setPreviewActive(value: boolean) { + if (this.preview.setPreviewActive(value)) { + this.changed.dispatch(); + } + } + + cachePreview(response: ReturnType) { + this.preview.cachePreview(response); + this.changed.dispatch(); + } + toJSON() { const { focusSegment, sinks, sources } = this; return { @@ -1628,6 +1674,70 @@ class GraphConnection extends SegmentationGraphSourceConnection { } } + private showSplitPreviewMessages() { + const { preview } = this.state.multicutState; + if (!preview.previewActive || !preview.hasCachedPreview) return; + if (preview.isSplitIllegal) { + StatusMessage.showTemporaryMessage( + "This split preview is illegal and cannot be submitted until the multicut seeds change.", + 7000, + ); + } + if (preview.connectedComponents.length > 2) { + StatusMessage.showTemporaryMessage( + "Additional split preview components are shown with the default preview color.", + 7000, + ); + } + } + + async requestSplitPreview(annotationToNanometers: Float64Array) { + const { + state: { multicutState }, + } = this; + const { sinks, sources, preview } = multicutState; + if (sinks.size === 0 || sources.size === 0) { + StatusMessage.showTemporaryMessage( + "You must select at least one red and one blue supervoxel to preview a split.", + 7000, + ); + return false; + } + if (preview.previewPending) { + return false; + } + if (preview.hasCachedPreview) { + multicutState.setPreviewActive(true); + this.showSplitPreviewMessages(); + return true; + } + multicutState.setPreviewPending(true); + try { + const response = await this.graph.graphServer.splitPreview( + [...sinks].map((x) => selectionInNanometers(x, annotationToNanometers)), + [...sources].map((x) => + selectionInNanometers(x, annotationToNanometers), + ), + ); + multicutState.cachePreview(response); + this.showSplitPreviewMessages(); + return true; + } catch { + multicutState.setPreviewPending(false); + return false; + } + } + + applyMulticutPreviewDisplay() { + return applySplitPreviewToTemporaryState( + this.segmentsState, + this.state.multicutState.preview.connectedComponents, + this.state.multicutState.focusSegment.value === undefined + ? [] + : [this.state.multicutState.focusSegment.value], + ); + } + async submitMulticut(annotationToNanometers: Float64Array): Promise { const { state: { multicutState }, @@ -1944,6 +2054,28 @@ class GrapheneGraphServerInterface { return final; } + async splitPreview(first: SegmentSelection[], second: SegmentSelection[]) { + const { fetchOkImpl, baseUrl } = this.httpSource; + const promise = fetchOkImpl( + `${baseUrl}/graph/split_preview?int64_as_str=1`, + { + method: "POST", + body: JSON.stringify({ + sources: first.map((x) => [String(x.segmentId), ...x.position]), + sinks: second.map((x) => [String(x.segmentId), ...x.position]), + }), + }, + ); + const jsonResp = await withErrorMessageHTTP( + promise.then((response) => response.json()), + { + initialMessage: `Calculating split preview for ${first.length + second.length} seeds`, + errorPrefix: "Split preview failed: ", + }, + ); + return parseGrapheneSplitPreviewResponse(jsonResp); + } + async filterLatestRoots(segments: bigint[]): Promise { const { fetchOkImpl, baseUrl } = this.httpSource; const url = `${baseUrl}/is_latest_roots`; @@ -2515,11 +2647,23 @@ class MulticutSegmentsTool extends LayerTool { makeToolActivationStatusMessageWithHeader(activation); header.textContent = "Multicut segments"; body.classList.add("graphene-tool-status", "graphene-multicut"); + const blockPreviewEdit = (message: string) => { + if (!multicutState.preview.previewActive) return false; + StatusMessage.showTemporaryMessage(message, 5000); + return true; + }; body.appendChild( makeIcon({ text: "Swap", title: "Swap group", onClick: () => { + if ( + blockPreviewEdit( + "Exit split preview before changing the multicut groups.", + ) + ) { + return; + } multicutState.swapGroup(); }, }), @@ -2533,13 +2677,43 @@ class MulticutSegmentsTool extends LayerTool { }, }), ); + const getAnnotationToNanometers = () => + getGraphLoadedSubsource( + this.layer, + )!.loadedDataSource.transform.inputSpace.value.scales.map( + (x) => x / 1e-9, + ); + const previewAction = async () => { + if (multicutState.preview.previewPending) { + return; + } + if (multicutState.preview.previewActive) { + multicutState.setPreviewActive(false); + return; + } + await graphConnection.requestSplitPreview(getAnnotationToNanometers()); + }; + const previewIcon = makeIcon({ + text: "Preview", + title: "Split preview", + onClick: () => { + previewAction(); + }, + }); + body.appendChild(previewIcon); const submitAction = async () => { - submitIcon.classList.toggle("disabled", true); - const loadedSubsource = getGraphLoadedSubsource(this.layer)!; - const annotationToNanometers = - loadedSubsource.loadedDataSource.transform.inputSpace.value.scales.map( - (x) => x / 1e-9, + if ( + multicutState.preview.hasCachedPreview && + multicutState.preview.isSplitIllegal + ) { + StatusMessage.showTemporaryMessage( + "This split preview is illegal. Change the multicut seeds before submitting.", + 7000, ); + return; + } + submitIcon.classList.toggle("disabled", true); + const annotationToNanometers = getAnnotationToNanometers(); graphConnection.submitMulticut(annotationToNanometers).then((success) => { submitIcon.classList.toggle("disabled", false); if (success) { @@ -2569,6 +2743,7 @@ class MulticutSegmentsTool extends LayerTool { activation.bindInputEventMap(MULTICUT_SEGMENTS_INPUT_EVENT_MAP); activation.registerDisposer(() => { + multicutState.setPreviewActive(false); resetMulticutDisplay(); displayState.baseSegmentHighlighting.value = priorBaseSegmentHighlighting; displayState.highlightColor.value = priorHighlightColor; @@ -2580,8 +2755,29 @@ class MulticutSegmentsTool extends LayerTool { displayState.tempSegmentDefaultColor2d.value = undefined; displayState.highlightColor.value = undefined; }; + const updatePreviewIcon = () => { + previewIcon.classList.toggle( + "disabled", + multicutState.preview.previewPending, + ); + previewIcon.title = multicutState.preview.previewActive + ? "Turn off split preview" + : multicutState.preview.hasCachedPreview + ? "Show cached split preview" + : "Split preview"; + previewIcon.style.borderStyle = multicutState.preview.previewActive + ? "inset" + : ""; + previewIcon.style.filter = multicutState.preview.previewActive + ? "invert(0.15)" + : ""; + previewIcon.style.webkitFilter = multicutState.preview.previewActive + ? "invert(0.15)" + : ""; + }; const updateMulticutDisplay = () => { resetMulticutDisplay(); + updatePreviewIcon(); activeGroupIndicator.classList.toggle( "blueGroup", multicutState.blueGroup.value, @@ -2589,6 +2785,28 @@ class MulticutSegmentsTool extends LayerTool { const focusSegment = multicutState.focusSegment.value; if (focusSegment === undefined) return; displayState.baseSegmentHighlighting.value = true; + if ( + multicutState.preview.previewActive && + multicutState.preview.hasCachedPreview + ) { + const previewRepresentatives = + graphConnection.applyMulticutPreviewDisplay(); + displayState.tempSegmentDefaultColor2d.value = MULTICUT_OFF_COLOR; + if (previewRepresentatives[0] !== undefined) { + displayState.tempSegmentStatedColors2d.value.set( + previewRepresentatives[0], + RED_COLOR_SEGMENT_PACKED, + ); + } + if (previewRepresentatives[1] !== undefined) { + displayState.tempSegmentStatedColors2d.value.set( + previewRepresentatives[1], + BLUE_COLOR_SEGMENT_PACKED, + ); + } + displayState.useTempSegmentStatedColors2d.value = true; + return; + } displayState.highlightColor.value = multicutState.blueGroup.value ? BLUE_COLOR_HIGHTLIGHT : RED_COLOR_HIGHLIGHT; @@ -2642,10 +2860,24 @@ class MulticutSegmentsTool extends LayerTool { ); activation.bindAction("swap-group", (event) => { event.stopPropagation(); + if ( + blockPreviewEdit( + "Exit split preview before changing the multicut groups.", + ) + ) { + return; + } multicutState.swapGroup(); }); activation.bindAction("set-anchor", (event) => { event.stopPropagation(); + if ( + blockPreviewEdit( + "Exit split preview before changing the multicut seeds.", + ) + ) { + return; + } const currentSegmentSelection = maybeGetSelection( this, segmentationGroupState.visibleSegments, diff --git a/src/datasource/graphene/split_preview.spec.ts b/src/datasource/graphene/split_preview.spec.ts new file mode 100644 index 0000000000..2566a90ae8 --- /dev/null +++ b/src/datasource/graphene/split_preview.spec.ts @@ -0,0 +1,100 @@ +/** + * @license + * Copyright 2016 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { describe, expect, it } from "vitest"; +import { + applySplitPreviewToTemporaryState, + MulticutSplitPreviewState, + parseGrapheneSplitPreviewResponse, +} from "#src/datasource/graphene/split_preview.js"; +import { VisibleSegmentEquivalencePolicy } from "#src/segmentation_graph/segment_id.js"; +import { SharedDisjointUint64Sets } from "#src/shared_disjoint_sets.js"; +import { WatchableValue } from "#src/trackable_value.js"; +import { Uint64Set } from "#src/uint64_set.js"; + +describe("graphene split preview helpers", () => { + it("parses split preview responses", () => { + expect( + parseGrapheneSplitPreviewResponse({ + supervoxel_connected_components: [["1", "2"], ["9"]], + illegal_split: true, + }), + ).toEqual({ + connectedComponents: [[1n, 2n], [9n]], + isSplitIllegal: true, + }); + }); + + it("invalidates cached multicut preview state", () => { + const state = new MulticutSplitPreviewState(); + expect(state.invalidate()).toBe(false); + expect(state.setPending(true)).toBe(true); + state.cachePreview({ + connectedComponents: [[1n, 2n], [3n]], + isSplitIllegal: true, + }); + expect(state.previewPending).toBe(false); + expect(state.previewActive).toBe(true); + expect(state.hasCachedPreview).toBe(true); + expect(state.isSplitIllegal).toBe(true); + expect(state.setPreviewActive(false)).toBe(true); + expect(state.previewActive).toBe(false); + expect(state.setPreviewActive(true)).toBe(true); + expect(state.previewActive).toBe(true); + expect(state.invalidate()).toBe(true); + expect(state.connectedComponents).toEqual([]); + expect(state.isSplitIllegal).toBe(false); + expect(state.previewActive).toBe(false); + expect(state.previewPending).toBe(false); + expect(state.hasCachedPreview).toBe(false); + }); + + it("applies connected components to temporary visible segments", () => { + const temporaryVisibleSegments = new Uint64Set(); + temporaryVisibleSegments.add(100n); + const temporarySegmentEquivalences = new SharedDisjointUint64Sets(); + temporarySegmentEquivalences.disjointSets.visibleSegmentEquivalencePolicy = + new WatchableValue( + VisibleSegmentEquivalencePolicy.MAX_REPRESENTATIVE | + VisibleSegmentEquivalencePolicy.NONREPRESENTATIVE_EXCLUDED, + ); + temporarySegmentEquivalences.link(10n, 11n); + const state = { + temporaryVisibleSegments, + temporarySegmentEquivalences, + useTemporaryVisibleSegments: new WatchableValue(false), + useTemporarySegmentEquivalences: new WatchableValue(false), + } as any; + const representatives = applySplitPreviewToTemporaryState( + state, + [[1n, 3n, 2n], [4n, 5n], [9n]], + [42n], + ); + expect(representatives).toEqual([3n, 5n, 9n]); + expect(state.useTemporaryVisibleSegments.value).toBe(true); + expect(state.useTemporarySegmentEquivalences.value).toBe(true); + expect(state.temporaryVisibleSegments.toJSON()).toEqual([ + "3", + "42", + "5", + "9", + ]); + expect(state.temporarySegmentEquivalences.get(1n)).toBe(3n); + expect(state.temporarySegmentEquivalences.get(2n)).toBe(3n); + expect(state.temporarySegmentEquivalences.get(4n)).toBe(5n); + expect(state.temporarySegmentEquivalences.get(9n)).toBe(9n); + }); +}); diff --git a/src/datasource/graphene/split_preview.ts b/src/datasource/graphene/split_preview.ts new file mode 100644 index 0000000000..c7c6ef76cb --- /dev/null +++ b/src/datasource/graphene/split_preview.ts @@ -0,0 +1,119 @@ +/** + * @license + * Copyright 2016 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { VisibleSegmentsState } from "#src/segmentation_display_state/base.js"; +import { + parseArray, + parseUint64, + verifyBoolean, + verifyObjectProperty, +} from "#src/util/json.js"; + +export interface GrapheneSplitPreview { + connectedComponents: bigint[][]; + isSplitIllegal: boolean; +} + +export function parseGrapheneSplitPreviewResponse( + response: unknown, +): GrapheneSplitPreview { + const connectedComponents = verifyObjectProperty( + response, + "supervoxel_connected_components", + (value) => + parseArray(value, (component) => parseArray(component, parseUint64)), + ); + const isSplitIllegal = verifyObjectProperty( + response, + "illegal_split", + verifyBoolean, + ); + return { + connectedComponents, + isSplitIllegal, + }; +} + +export class MulticutSplitPreviewState { + connectedComponents: bigint[][] = []; + isSplitIllegal = false; + previewActive = false; + previewPending = false; + hasCachedPreview = false; + + invalidate() { + const changed = + this.previewActive || + this.previewPending || + this.hasCachedPreview || + this.isSplitIllegal || + this.connectedComponents.length !== 0; + this.connectedComponents = []; + this.isSplitIllegal = false; + this.previewActive = false; + this.previewPending = false; + this.hasCachedPreview = false; + return changed; + } + + setPending(value: boolean) { + if (this.previewPending === value) return false; + this.previewPending = value; + return true; + } + + setPreviewActive(value: boolean) { + if (value && !this.hasCachedPreview) return false; + if (this.previewActive === value) return false; + this.previewActive = value; + return true; + } + + cachePreview(preview: GrapheneSplitPreview) { + this.connectedComponents = preview.connectedComponents.map((component) => [ + ...component, + ]); + this.isSplitIllegal = preview.isSplitIllegal; + this.hasCachedPreview = true; + this.previewPending = false; + this.previewActive = true; + } +} + +export function applySplitPreviewToTemporaryState( + state: VisibleSegmentsState, + connectedComponents: readonly (readonly bigint[])[], + preservedVisibleSegments: readonly bigint[] = [], +) { + const representatives: bigint[] = []; + const tempVisibleSegments = state.temporaryVisibleSegments; + const tempEquivalences = state.temporarySegmentEquivalences; + tempVisibleSegments.clear(); + tempEquivalences.clear(); + state.useTemporaryVisibleSegments.value = true; + state.useTemporarySegmentEquivalences.value = true; + for (const segment of preservedVisibleSegments) { + tempVisibleSegments.add(segment); + } + for (const component of connectedComponents) { + if (component.length === 0) continue; + tempEquivalences.linkAll([...component]); + const representative = tempEquivalences.get(component[0]); + tempVisibleSegments.add(representative); + representatives.push(representative); + } + return representatives; +}