Skip to content

feat(segmentation_user_layer): add ability to color segmentation based on segment properties #852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
+285 −19
Draft

feat(segmentation_user_layer): add ability to color segmentation based on segment properties #852

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
234 changes: 232 additions & 2 deletions src/layer/segmentation/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,22 +119,56 @@ import {
} from "#src/util/color.js";
import type { Borrowed, Owned } from "#src/util/disposable.js";
import { RefCounted } from "#src/util/disposable.js";
import type { vec3, vec4 } from "#src/util/geom.js";
import type { vec4 } from "#src/util/geom.js";
import { vec3 } from "#src/util/geom.js";
import {
parseArray,
parseUint64,
verifyArray,
verifyBoolean,
verifyFiniteNonNegativeFloat,
verifyFloat,
verifyObject,
verifyObjectAsMap,
verifyObjectProperty,
verifyOptionalObjectProperty,
verifyString,
} from "#src/util/json.js";
import { computeInvlerp } from "#src/util/lerp.js";
import { Signal } from "#src/util/signal.js";
import { makeWatchableShaderError } from "#src/webgl/dynamic_shader.js";
import type { DependentViewContext } from "#src/widget/dependent_view_widget.js";
import { registerLayerShaderControlsTool } from "#src/widget/shader_controls.js";

const MAX_LAYER_BAR_UI_INDICATOR_COLORS = 6;

export interface SegmentPropertyColorBase {
type: "tag" | "numeric";
active: boolean;
}

export interface SegmentPropertyColorNumeric extends SegmentPropertyColorBase {
type: "numeric";
active: boolean;
property: string;
options: {
min?: number;
max?: number;
minColor: string;
maxColor: string;
};
}

export interface SegmentPropertyColorTag extends SegmentPropertyColorBase {
type: "tag";
active: boolean;
map: Map<string, string>;
}

export type SegmentPropertyColor =
| SegmentPropertyColorNumeric
| SegmentPropertyColorTag;

export class SegmentationUserLayerGroupState
extends RefCounted
implements SegmentationGroupState
Expand Down Expand Up @@ -298,9 +332,84 @@ export class SegmentationUserLayerColorGroupState
specificationChanged = new Signal();
constructor(public layer: SegmentationUserLayer) {
super();

this.segmentPropertyColorsMap = this.registerDisposer(
makeCachedDerivedWatchableValue(
(segmentPropertyMap, segmentPropertyColors) => {
const map = new Uint64Map();
if (!segmentPropertyMap) {
return map;
}
const { tags: tagsProperty } = segmentPropertyMap;
for (const propertyColor of segmentPropertyColors) {
if (!propertyColor.active) continue;
if (propertyColor.type === "tag" && tagsProperty) {
const { tags, values } = tagsProperty;
const bigIntColors = new Map<string, bigint>();
for (const [tag, colorString] of propertyColor.map!.entries()) {
const color = parseRGBColorSpecification(colorString);
bigIntColors.set(tag, BigInt(packColor(color)));
}
for (const [index, id] of (
segmentPropertyMap.segmentPropertyMap.inlineProperties?.ids ||
[]
).entries()) {
if (map.has(id)) continue; // prioritize first color match
const tagIndices = values[index];
const segmentTags = new Set(
tagIndices.split("").map((x) => tags[x.charCodeAt(0)]),
);
for (const [tag, color] of bigIntColors.entries()) {
if (segmentTags.has(tag)) {
map.set(id, color);
break;
}
}
}
} else if (propertyColor.type === "numeric") {
const options = propertyColor.options;
const { numericalProperties } = segmentPropertyMap;
const numericalProperty = numericalProperties.find(
(x) => x.id === propertyColor.property,
);
if (numericalProperty) {
const minColor = parseRGBColorSpecification(options.minColor);
const maxColor = parseRGBColorSpecification(options.maxColor);
const color = vec3.create();
const { values, bounds } = numericalProperty;
const min = Number(options.min ?? bounds[0]);
const max = Number(options.max ?? bounds[1]);
for (const [index, id] of (
segmentPropertyMap.segmentPropertyMap.inlineProperties?.ids ||
[]
).entries()) {
if (map.has(id)) continue; // prioritize first color match
const value = values[index];
if (Number.isNaN(value)) continue;
const normalized = Math.min(
1,
Math.max(0, computeInvlerp([min, max], value)),
);
vec3.lerp(color, minColor, maxColor, normalized);
map.set(id, BigInt(packColor(color)));
}
}
}
}
return map;
},
[
this.layer.displayState.originalSegmentationGroupState
.segmentPropertyMap,
this.segmentPropertyColors,
],
),
);

const { specificationChanged } = this;
this.segmentColorHash.changed.add(specificationChanged.dispatch);
this.segmentStatedColors.changed.add(specificationChanged.dispatch);
this.segmentPropertyColorsMap.changed.add(specificationChanged.dispatch);
this.tempSegmentStatedColors2d.changed.add(specificationChanged.dispatch);
this.segmentDefaultColor.changed.add(specificationChanged.dispatch);
this.tempSegmentDefaultColor2d.changed.add(specificationChanged.dispatch);
Expand Down Expand Up @@ -332,33 +441,140 @@ export class SegmentationUserLayerColorGroupState
}
},
);
verifyOptionalObjectProperty(
specification,
json_keys.SEGMENT_PROPERTY_COLORS_JSON_KEY,
(value) => {
const segmentPropertyColors = verifyArray(value);
for (const propertyColor of segmentPropertyColors) {
verifyObject(propertyColor);
const type = verifyObjectProperty(
propertyColor,
"type",
verifyString,
);
const active = verifyOptionalObjectProperty(
propertyColor,
"active",
verifyBoolean,
);
if (type === "tag") {
const map = verifyObjectProperty(propertyColor, "map", (x) => {
verifyObject(x);
const res = new Map<string, string>();
for (const [tag, value] of Object.entries(x)) {
const color = verifyString(value);
res.set(tag, color);
}
return res;
});
this.segmentPropertyColors.value.push({
type,
active: active ?? true,
map,
});
} else if (type === "numeric") {
const property = verifyObjectProperty(
propertyColor,
"property",
verifyString,
);
const options = verifyObjectProperty(
propertyColor,
"options",
(x) => {
verifyObject(x);
const min = verifyOptionalObjectProperty(x, "min", verifyFloat);
const max = verifyOptionalObjectProperty(x, "max", verifyFloat);
const minColor = verifyObjectProperty(
x,
"minColor",
verifyString,
);
const maxColor = verifyObjectProperty(
x,
"maxColor",
verifyString,
);
return {
min,
max,
minColor,
maxColor,
};
},
);
this.segmentPropertyColors.value.push({
type,
property,
active: active ?? true,
options,
});
} else {
throw new Error(`unsupported type: ${type}`);
}
}
this.segmentPropertyColors.changed.dispatch();
},
);
}

toJSON() {
const x: any = {};
x[json_keys.COLOR_SEED_JSON_KEY] = this.segmentColorHash.toJSON();
x[json_keys.SEGMENT_DEFAULT_COLOR_JSON_KEY] =
this.segmentDefaultColor.toJSON();
const { segmentStatedColors } = this;
const {
segmentStatedColors,
segmentPropertyColors: { value: segmentPropertyColors },
} = this;
if (segmentStatedColors.size > 0) {
const j: any = (x[json_keys.SEGMENT_STATED_COLORS_JSON_KEY] = {});
for (const [key, value] of segmentStatedColors) {
j[key.toString()] = serializeColor(unpackRGB(Number(value)));
}
}
if (segmentPropertyColors.length > 0) {
x[json_keys.SEGMENT_PROPERTY_COLORS_JSON_KEY] = segmentPropertyColors.map(
(propertyColor) => {
const { type, active } = propertyColor;
if (type === "numeric") {
const { property, options } = propertyColor;
return {
type,
active,
property,
options,
};
} else {
const { map } = propertyColor;
return {
type,
active,
map: Object.fromEntries(
map.entries().map(([k, v]) => [k.toString(), v]),
),
};
}
},
);
}
return x;
}

assignFrom(other: SegmentationUserLayerColorGroupState) {
this.segmentColorHash.value = other.segmentColorHash.value;
this.segmentStatedColors.assignFrom(other.segmentStatedColors);
this.segmentPropertyColors.value = other.segmentPropertyColors.value;
this.tempSegmentStatedColors2d.assignFrom(other.tempSegmentStatedColors2d);
this.segmentDefaultColor.value = other.segmentDefaultColor.value;
this.highlightColor.value = other.highlightColor.value;
}

segmentColorHash = SegmentColorHash.getDefault();
segmentStatedColors = this.registerDisposer(new Uint64Map());
segmentPropertyColors = new WatchableValue<SegmentPropertyColor[]>([]);
segmentPropertyColorsMap: WatchableValueInterface<Uint64Map>;
tempSegmentStatedColors2d = this.registerDisposer(new Uint64Map());
segmentDefaultColor = new TrackableOptionalRGB();
tempSegmentDefaultColor2d = new WatchableValue<vec3 | vec4 | undefined>(
Expand Down Expand Up @@ -484,6 +700,18 @@ class SegmentationUserLayerDisplayState implements SegmentationDisplayState {
(group) => group.segmentStatedColors,
),
);
this.segmentPropertyColors = this.layer.registerDisposer(
new IndirectTrackableValue(
this.segmentationColorGroupState,
(group) => group.segmentPropertyColors,
),
);
this.segmentPropertyColorsMap = this.layer.registerDisposer(
new IndirectTrackableValue(
this.segmentationColorGroupState,
(group) => group.segmentPropertyColorsMap,
),
);
this.tempSegmentStatedColors2d = this.layer.registerDisposer(
new IndirectTrackableValue(
this.segmentationColorGroupState,
Expand Down Expand Up @@ -564,6 +792,8 @@ class SegmentationUserLayerDisplayState implements SegmentationDisplayState {
segmentColorHash: TrackableValueInterface<number>;
segmentStatedColors: WatchableValueInterface<Uint64Map>;
tempSegmentStatedColors2d: WatchableValueInterface<Uint64Map>;
segmentPropertyColors: WatchableValueInterface<SegmentPropertyColor[]>;
segmentPropertyColorsMap: WatchableValueInterface<Uint64Map>;
segmentDefaultColor: WatchableValueInterface<vec3 | undefined>;
tempSegmentDefaultColor2d: WatchableValueInterface<vec3 | vec4 | undefined>;
highlightColor: WatchableValueInterface<vec4 | undefined>;
Expand Down
1 change: 1 addition & 0 deletions src/layer/segmentation/json_keys.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ export const SEGMENTS_JSON_KEY = "segments";
export const EQUIVALENCES_JSON_KEY = "equivalences";
export const COLOR_SEED_JSON_KEY = "colorSeed";
export const SEGMENT_STATED_COLORS_JSON_KEY = "segmentColors";
export const SEGMENT_PROPERTY_COLORS_JSON_KEY = "segmentPropertyColors";
export const MESH_RENDER_SCALE_JSON_KEY = "meshRenderScale";
export const CROSS_SECTION_RENDER_SCALE_JSON_KEY = "crossSectionRenderScale";
export const SKELETON_RENDERING_JSON_KEY = "skeletonRendering";
Expand Down
23 changes: 21 additions & 2 deletions src/segmentation_display_state/frontend.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@ import type { LayerChunkProgressInfo } from "#src/chunk_manager/base.js";
import type { ChunkManager } from "#src/chunk_manager/frontend.js";
import { ChunkRenderLayerFrontend } from "#src/chunk_manager/frontend.js";
import type { LayerSelectedValues } from "#src/layer/index.js";
import type { SegmentationUserLayer } from "#src/layer/segmentation/index.js";
import type {
SegmentationUserLayer,
SegmentPropertyColor,
} from "#src/layer/segmentation/index.js";
import type { PickIDManager } from "#src/object_picking.js";
import type { WatchableRenderLayerTransform } from "#src/render_coordinate_transform.js";
import type { RenderScaleHistogram } from "#src/render_scale_statistics.js";
Expand Down Expand Up @@ -178,6 +181,8 @@ export interface SegmentationColorGroupState {
segmentColorHash: SegmentColorHash;
segmentStatedColors: Uint64Map;
tempSegmentStatedColors2d: Uint64Map;
segmentPropertyColors: WatchableValueInterface<SegmentPropertyColor[]>;
segmentPropertyColorsMap: WatchableValueInterface<Uint64Map>;
segmentDefaultColor: WatchableValueInterface<vec3 | undefined>;
tempSegmentDefaultColor2d: WatchableValueInterface<vec3 | vec4 | undefined>;
}
Expand All @@ -199,6 +204,7 @@ export interface SegmentationDisplayState {
hideSegmentZero: WatchableValueInterface<boolean>;
segmentColorHash: WatchableValueInterface<number>;
segmentStatedColors: WatchableValueInterface<Uint64Map>;
segmentPropertyColorsMap: WatchableValueInterface<Uint64Map>;
tempSegmentStatedColors2d: WatchableValueInterface<Uint64Map>;
useTempSegmentStatedColors2d: WatchableValueInterface<boolean>;
segmentDefaultColor: WatchableValueInterface<vec3 | undefined>;
Expand Down Expand Up @@ -934,7 +940,10 @@ export function getBaseObjectColor(
return color;
}
const colorGroupState = displayState.segmentationColorGroupState.value;
const { segmentStatedColors } = colorGroupState;
const {
segmentStatedColors,
segmentPropertyColorsMap: { value: segmentPropertyColorsMap },
} = colorGroupState;
let statedColor: bigint | undefined;
if (
segmentStatedColors.size !== 0 &&
Expand All @@ -947,6 +956,16 @@ export function getBaseObjectColor(
color[2] = (Number(statedColor & 0xff0000n) >>> 16) / 255.0;
return color;
}
let propertyColor: bigint | undefined;
if (
segmentPropertyColorsMap.size !== 0 &&
(propertyColor = segmentPropertyColorsMap.get(objectId)) !== undefined
) {
color[0] = Number(propertyColor & 0x0000ffn) / 255.0;
color[1] = (Number(propertyColor & 0x00ff00n) >>> 8) / 255.0;
color[2] = (Number(propertyColor & 0xff0000n) >>> 16) / 255.0;
return color;
}
const segmentDefaultColor = colorGroupState.segmentDefaultColor.value;
if (segmentDefaultColor !== undefined) {
color[0] = segmentDefaultColor[0];
Expand Down
Loading
Loading