diff --git a/app/packages/looker/package.json b/app/packages/looker/package.json index 2f41057ad8..f42a7eec9c 100644 --- a/app/packages/looker/package.json +++ b/app/packages/looker/package.json @@ -18,21 +18,27 @@ "README.md" ], "dependencies": { + "@jimp/tiff": "^0.22.12", "@ungap/event-target": "^0.2.2", "@xmldom/xmldom": "^0.8.6", "copy-to-clipboard": "^3.3.1", "decode-tiff": "^0.2.1", + "geotiff": "^2.1.3", "immutable": "^4.0.0-rc.12", "lodash": "^4.17.21", "lru-cache": "^11.0.1", "mime": "^2.5.2", "monotone-convex-hull-2d": "^1.0.1", + "tiff": "^6.1.1", + "utif": "^3.1.0", + "utif2": "^4.1.0", "uuid": "^8.3.2" }, "devDependencies": { "@rollup/plugin-inject": "^5.0.2", "@types/color-string": "^1.5.0", "@types/lru-cache": "^7.10.10", + "@types/utif": "^3.0.5", "@types/uuid": "^8.3.0", "buffer": "^6.0.3", "prettier": "^2.7.1", diff --git a/app/packages/looker/src/state.ts b/app/packages/looker/src/state.ts index 996eb23ef9..84c9af8b48 100644 --- a/app/packages/looker/src/state.ts +++ b/app/packages/looker/src/state.ts @@ -2,11 +2,15 @@ * Copyright 2017-2024, Voxel51, Inc. */ -import { BufferManager } from "@fiftyone/utilities"; -import { ImaVidFramesController } from "./lookers/imavid/controller"; -import { Overlay } from "./overlays/base"; - -import { AppError, COLOR_BY, Schema, Stage } from "@fiftyone/utilities"; +import type { + AppError, + BufferManager, + COLOR_BY, + Schema, + Stage, +} from "@fiftyone/utilities"; +import type { ImaVidFramesController } from "./lookers/imavid/controller"; +import type { Overlay } from "./overlays/base"; export type Optional = { [P in keyof T]?: Optional; @@ -70,27 +74,6 @@ export type OrthogrpahicProjectionMetadata = { normal: [number, number, number]; }; -export type GenericLabel = { - [labelKey: string]: { - [field: string]: unknown; - }; - // todo: add other label types -}; - -export type Sample = { - metadata: { - width: number; - height: number; - mime_type?: string; - }; - _id: string; - id: string; - filepath: string; - tags: string[]; - _label_tags: string[]; - _media_type: "image" | "video" | "point-cloud" | "3d"; -} & GenericLabel; - export interface LabelData { labelId: string; field: string; @@ -127,8 +110,8 @@ export type Action = ( ) => void; export enum ControlEventKeyType { - HOLD, - KEY_DOWN, + HOLD = "HOLD", + KEY_DOWN = "KEY_DOWN", } export interface Control { eventKeys?: string | string[]; diff --git a/app/packages/looker/src/worker/decoders/canvas.ts b/app/packages/looker/src/worker/decoders/canvas.ts index 93ecd6c0ba..439279fb01 100644 --- a/app/packages/looker/src/worker/decoders/canvas.ts +++ b/app/packages/looker/src/worker/decoders/canvas.ts @@ -1,5 +1,4 @@ -import type { OverlayMask } from "../../numpy"; -import { enqueueFetch } from "../pooled-fetch"; +import type { OverlayMask } from "./types"; const PNG_SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]; /** @@ -34,13 +33,7 @@ const getPngcolorType = async (blob: Blob): Promise => { return colorType; }; -const decodeWithCanvas = async (url: string): Promise => { - const overlayImageFetchResponse = await enqueueFetch({ - url, - options: { priority: "low" }, - }); - const blob = await overlayImageFetchResponse.blob(); - +const decodeWithCanvas = async (blob: Blob): Promise => { let channels = 4; if (blob.type === "image/png") { diff --git a/app/packages/looker/src/worker/decoders/index.ts b/app/packages/looker/src/worker/decoders/index.ts index 21076d98a0..f96e21c3d2 100644 --- a/app/packages/looker/src/worker/decoders/index.ts +++ b/app/packages/looker/src/worker/decoders/index.ts @@ -1,37 +1,13 @@ import { getSampleSrc } from "@fiftyone/state/src/recoil/utils"; import { DETECTION, DETECTIONS } from "@fiftyone/utilities"; import type { Coloring, CustomizeColor } from "../.."; -import type { OverlayMask } from "../../numpy"; import type { Colorscale } from "../../state"; +import { enqueueFetch } from "../pooled-fetch"; import { getOverlayFieldFromCls } from "../shared"; import decodeCanvas from "./canvas"; +import decodeInline from "./inline"; import decodeTiff from "./tiff"; - -type Decoder = (url: string) => Promise; - -const IMAGE_DECODERS: { [key: string]: Decoder } = { - jpeg: decodeCanvas, - jpg: decodeCanvas, - png: decodeCanvas, - tif: decodeTiff, - tiff: decodeTiff, -}; - -async function decodeImage(url: string) { - const extension = url.split(".").slice(-1)[0]; - return IMAGE_DECODERS[extension](url); -} - -export type IntermediateMask = { - data: OverlayMask; - image: ArrayBuffer; -}; - -interface Decodeable { - mask_path?: string; - map_path?: string; - detections?: { mask_path?: string }[]; -} +import type { BlobDecoder, Decodeables } from "./types"; interface DecodeParameters { cls: string; @@ -39,13 +15,21 @@ interface DecodeParameters { colorscale: Colorscale; customizeColorSetting: CustomizeColor[]; field: string; - label: Decodeable; + label: Decodeables; maskPathDecodingPromises?: Promise[]; maskTargetsBuffers?: ArrayBuffer[]; overlayCollectionProcessingParams?: { idx: number; cls: string }; sources: { [path: string]: string }; } +const IMAGE_DECODERS: { [key: string]: BlobDecoder } = { + jpeg: decodeCanvas, + jpg: decodeCanvas, + png: decodeCanvas, + tif: decodeTiff, + tiff: decodeTiff, +}; + /** * Some label types (example: segmentation, heatmap) can have their overlay * data stored on-disk, we want to impute the relevant mask property of these @@ -74,21 +58,31 @@ const decode = async ({ field, label, ...params }: DecodeParameters) => { const overlayPathField = overlayFields.disk; const overlayField = overlayFields.canonical; + const data = label[overlayField]; + if (typeof data === "string") { + const intermediate = decodeInline(data); + if (!intermediate) { + return; + } + + params.maskTargetsBuffers.push(intermediate.data.buffer); + label[overlayField] = intermediate; + return; + } + if (!Object.hasOwn(label, overlayPathField)) { return; } - // it's possible we're just re-coloring, in which case re-init mask image - // and set bitmap to null - if (label[overlayField]?.bitmap && !label[overlayField]?.image) { + if (data?.bitmap && !data?.image) { const height = label[overlayField].bitmap.height; const width = label[overlayField].bitmap.width; // close the copied bitmap - label[overlayField].bitmap.close(); - label[overlayField].bitmap = null; + data.bitmap.close(); + data.bitmap = null; - label[overlayField].image = new ArrayBuffer(height * width * 4); + data.image = new ArrayBuffer(height * width * 4); return; } @@ -110,7 +104,13 @@ const decode = async ({ field, label, ...params }: DecodeParameters) => { // convert absolute file path to a URL that we can "fetch" from const overlayImageUrl = getSampleSrc(source || label[overlayPathField]); - const overlayMask = await decodeImage(overlayImageUrl); + const overlayImageFetchResponse = await enqueueFetch({ + url: overlayImageUrl, + options: { priority: "low" }, + }); + const blob = await overlayImageFetchResponse.blob(); + const extension = overlayImageUrl.split(".").slice(-1)[0]; + const overlayMask = await IMAGE_DECODERS[extension](blob); const [overlayHeight, overlayWidth] = overlayMask.shape; // set the `mask` property for this label @@ -120,7 +120,7 @@ const decode = async ({ field, label, ...params }: DecodeParameters) => { label[overlayField] = { data: overlayMask, image: new ArrayBuffer(overlayWidth * overlayHeight * 4), - } as IntermediateMask; + }; // no need to transfer image's buffer // since we'll be constructing ImageBitmap and transfering that diff --git a/app/packages/looker/src/worker/decoders/inline.ts b/app/packages/looker/src/worker/decoders/inline.ts new file mode 100644 index 0000000000..74aabd3853 --- /dev/null +++ b/app/packages/looker/src/worker/decoders/inline.ts @@ -0,0 +1,23 @@ +import { deserialize } from "./numpy"; +import type { InlineDecodeable, IntermediateMask } from "./types"; + +export default (inline: InlineDecodeable): IntermediateMask | null => { + let base64: string; + if (typeof inline === "string") { + base64 = inline; + } else if (typeof inline?.$binary?.base64 === "string") { + base64 = inline.$binary.base64; + } + + if (!base64) { + return null; + } + + const data = deserialize(base64); + const [height, width] = data.shape; + + return { + data, + image: new ArrayBuffer(width * height * 4), + }; +}; diff --git a/app/packages/looker/src/numpy.ts b/app/packages/looker/src/worker/decoders/numpy.ts similarity index 80% rename from app/packages/looker/src/numpy.ts rename to app/packages/looker/src/worker/decoders/numpy.ts index f48c6abce1..94f8365cf5 100644 --- a/app/packages/looker/src/numpy.ts +++ b/app/packages/looker/src/worker/decoders/numpy.ts @@ -4,42 +4,7 @@ import { Buffer } from "buffer"; import pako from "./pako.js"; - -export { deserialize }; - -export const ARRAY_TYPES = { - Uint8Array, - Uint8ClampedArray, - Int8Array, - Uint16Array, - Int16Array, - Uint32Array, - Int32Array, - Float32Array, - Float64Array, - BigUint64Array, - BigInt64Array, -}; - -export type TypedArray = - | Uint8Array - | Uint8ClampedArray - | Int8Array - | Uint16Array - | Int16Array - | Uint32Array - | Int32Array - | Float32Array - | Float64Array - | BigUint64Array - | BigInt64Array; - -export interface OverlayMask { - buffer: ArrayBuffer; - shape: [number, number]; - channels: number; - arrayType: keyof typeof ARRAY_TYPES; -} +import type { OverlayMask } from "./types.js"; const DATA_TYPES = { // < = little-endian, > = big-endian, | = host architecture @@ -146,6 +111,6 @@ function parse(array: Uint8Array): OverlayMask { /** * Deserializes and parses a base64 encoded numpy array */ -function deserialize(compressedBase64Array: string): OverlayMask { +export function deserialize(compressedBase64Array: string): OverlayMask { return parse(pako.inflate(Buffer.from(compressedBase64Array, "base64"))); } diff --git a/app/packages/looker/src/pako.js b/app/packages/looker/src/worker/decoders/pako.js similarity index 100% rename from app/packages/looker/src/pako.js rename to app/packages/looker/src/worker/decoders/pako.js diff --git a/app/packages/looker/src/worker/decoders/tiff.ts b/app/packages/looker/src/worker/decoders/tiff.ts index d6e9fe5b77..ff67d1717d 100644 --- a/app/packages/looker/src/worker/decoders/tiff.ts +++ b/app/packages/looker/src/worker/decoders/tiff.ts @@ -1,14 +1,25 @@ -import { getFetchFunction } from "@fiftyone/utilities"; -import { decode } from "decode-tiff"; -import type { OverlayMask } from "../../numpy"; - -export default async (url: string) => { - const buffer: ArrayBuffer = await getFetchFunction()( - "GET", - url, - null, - "arrayBuffer" - ); - const result = decode(buffer); - return {} as OverlayMask; +import * as tiff from "geotiff"; +import type { OverlayMask } from "./types"; + +export default async (blob: Blob) => { + const r = await tiff.fromBlob(blob); + const image = await r.getImage(); + console.log(image.getHeight(), image.getWidth()); + console.log(r); + console.log(image); + + console.log(); + + const data = await image.readRGB(); + + if (!(data instanceof Uint8Array)) { + throw new Error("wrong"); + } + + return { + arrayType: "Uint8Array", + buffer: data.buffer, + shape: [4, 4], + channels: 2, + } as OverlayMask; }; diff --git a/app/packages/looker/src/worker/decoders/types.ts b/app/packages/looker/src/worker/decoders/types.ts new file mode 100644 index 0000000000..bd1bb7a5aa --- /dev/null +++ b/app/packages/looker/src/worker/decoders/types.ts @@ -0,0 +1,62 @@ +export type BlobDecoder = (url: Blob) => Promise; + +export const ARRAY_TYPES = { + Uint8Array, + Uint8ClampedArray, + Int8Array, + Uint16Array, + Int16Array, + Uint32Array, + Int32Array, + Float32Array, + Float64Array, + BigUint64Array, + BigInt64Array, +}; + +export type TypedArray = + | Uint8Array + | Uint8ClampedArray + | Int8Array + | Uint16Array + | Int16Array + | Uint32Array + | Int32Array + | Float32Array + | Float64Array + | BigUint64Array + | BigInt64Array; + +export type Decodeables = + | { + detections?: InstanceDecodeable[]; + } & InstanceDecodeable; + +export interface Decoded { + bitmap?: ImageBitmap; +} + +export type InlineDecodeable = + | { + $binary?: { base64: string }; + } + | string; + +interface InstanceDecodeable { + map?: Decoded & InlineDecodeable & IntermediateMask; + map_path?: string; + mask?: Decoded & InlineDecodeable & IntermediateMask; + mask_path?: string; +} + +export interface IntermediateMask { + data: OverlayMask; + image: ArrayBuffer; +} + +export interface OverlayMask { + buffer: ArrayBuffer; + shape: [number, number]; + channels: number; + arrayType: keyof typeof ARRAY_TYPES; +} diff --git a/app/packages/looker/src/worker/deserializer.test.ts b/app/packages/looker/src/worker/deserializer.test.ts deleted file mode 100644 index 837f28ae76..0000000000 --- a/app/packages/looker/src/worker/deserializer.test.ts +++ /dev/null @@ -1,29 +0,0 @@ -import { describe, expect, it } from "vitest"; - -import * as deserializer from "./deserializer"; - -describe("filter resolves correctly", () => { - it("skips undefined detection", () => { - expect( - deserializer.DeserializerFactory.Detection(undefined, []) - ).toBeUndefined(); - }); - - it("skips undefined detections", () => { - expect( - deserializer.DeserializerFactory.Detections(undefined, []) - ).toBeUndefined(); - }); - - it("skips undefined heatmap", () => { - expect( - deserializer.DeserializerFactory.Heatmap(undefined, []) - ).toBeUndefined(); - }); - - it("skips undefined segmenation", () => { - expect( - deserializer.DeserializerFactory.Segmentation(undefined, []) - ).toBeUndefined(); - }); -}); diff --git a/app/packages/looker/src/worker/deserializer.ts b/app/packages/looker/src/worker/deserializer.ts deleted file mode 100644 index a1ffc5e793..0000000000 --- a/app/packages/looker/src/worker/deserializer.ts +++ /dev/null @@ -1,68 +0,0 @@ -import { deserialize } from "../numpy"; - -const extractSerializedMask = ( - label: object, - maskProp: string -): string | undefined => { - if (typeof label?.[maskProp] === "string") { - return label[maskProp]; - } - - if (typeof label?.[maskProp]?.$binary?.base64 === "string") { - return label[maskProp].$binary.base64; - } - - return undefined; -}; - -export const DeserializerFactory = { - Detection: (label, buffers) => { - const serializedMask = extractSerializedMask(label, "mask"); - - if (serializedMask) { - const data = deserialize(serializedMask); - const [height, width] = data.shape; - label.mask = { - data, - image: new ArrayBuffer(width * height * 4), - }; - buffers.push(data.buffer); - } - }, - Detections: (labels, buffers) => { - const list = labels?.detections || []; - for (const label of list) { - DeserializerFactory.Detection(label, buffers); - } - }, - Heatmap: (label, buffers) => { - const serializedMask = extractSerializedMask(label, "map"); - - if (serializedMask) { - const data = deserialize(serializedMask); - const [height, width] = data.shape; - - label.map = { - data, - image: new ArrayBuffer(width * height * 4), - }; - - buffers.push(data.buffer); - } - }, - Segmentation: (label, buffers) => { - const serializedMask = extractSerializedMask(label, "mask"); - - if (serializedMask) { - const data = deserialize(serializedMask); - const [height, width] = data.shape; - - label.mask = { - data, - image: new ArrayBuffer(width * height * 4), - }; - - buffers.push(data.buffer); - } - }, -}; diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index a7846cd1c5..b1ff98e6c3 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -11,6 +11,7 @@ import { EMBEDDED_DOCUMENT, HEATMAP, LABEL_LIST, + PANOPTIC_SEGMENTATION, SEGMENTATION, VALID_LABEL_TYPES, getCls, @@ -29,7 +30,6 @@ import type { Sample, } from "../state"; import decodeImages from "./decoders"; -import { DeserializerFactory } from "./deserializer"; import { painter, resolveColor } from "./painters"; import type { ResolveColor, ResolveColorMethod } from "./painters/utils"; import { getOverlayFieldFromCls, mapId } from "./shared"; @@ -99,10 +99,6 @@ const processLabels = async ( ); } - if (cls in DeserializerFactory) { - DeserializerFactory[cls](label, maskTargetsBuffers); - } - if ([EMBEDDED_DOCUMENT, DYNAMIC_EMBEDDED_DOCUMENT].includes(cls)) { const [moreBitmapPromises, moreMaskTargetsBuffers] = await processLabels( @@ -162,6 +158,7 @@ const processLabels = async ( selectedLabelTags, }; let promise: Promise; + switch (cls) { case DETECTION: promise = painter.Detection(params); @@ -172,6 +169,9 @@ const processLabels = async ( case HEATMAP: promise = painter.Heatmap(params); break; + case PANOPTIC_SEGMENTATION: + promise = painter.PanopticSegmentation(params); + break; case SEGMENTATION: promise = painter.Segmentation(params); break; diff --git a/app/packages/looker/src/worker/painters/detection.ts b/app/packages/looker/src/worker/painters/detection.ts index 7e8dbd2794..e8b10a017b 100644 --- a/app/packages/looker/src/worker/painters/detection.ts +++ b/app/packages/looker/src/worker/painters/detection.ts @@ -1,8 +1,8 @@ import { COLOR_BY, get32BitColor } from "@fiftyone/utilities"; import colorString from "color-string"; -import { ARRAY_TYPES } from "../../numpy"; import type { DetectionLabel } from "../../overlays/detection"; import { getHashLabel, shouldShowLabelTag } from "../../overlays/util"; +import { ARRAY_TYPES } from "../decoders/types"; import type { Painter } from "./utils"; import { requestColor } from "./utils"; @@ -75,7 +75,7 @@ const detection: Painter = async ({ const valueColor = setting?.valueColors?.find((l) => { if (["none", "null", "undefined"].includes(l.value?.toLowerCase())) { return typeof label[key] === "string" - ? l.value?.toLowerCase === label[key] + ? l.value?.toLowerCase() === label[key] : !label[key]; } if (["True", "False"].includes(l.value?.toString())) { diff --git a/app/packages/looker/src/worker/painters/index.ts b/app/packages/looker/src/worker/painters/index.ts index 5a3f62b6a2..f78ab71aea 100644 --- a/app/packages/looker/src/worker/painters/index.ts +++ b/app/packages/looker/src/worker/painters/index.ts @@ -1,13 +1,15 @@ import type { DetectionLabel } from "../../overlays/detection"; +import type { IntermediateMask } from "../decoders/types"; import detection from "./detection"; import heatmap from "./heatmap"; +import panoptic from "./panoptic"; import segmentation from "./segmentation"; import type { Painter } from "./utils"; import { requestColor, resolveColor } from "./utils"; type Detections = (request: typeof requestColor) => ( params: Omit>[0], "label"> & { - label: { detections: DetectionLabel[] }; + label: { detections: (DetectionLabel & { mask: IntermediateMask })[] }; } ) => Promise; @@ -32,6 +34,7 @@ export const createPainter = (request: typeof requestColor) => ({ Detection: detection, Detections: detections(request), Heatmap: heatmap, + PanopticSegmentation: panoptic, Segmentation: segmentation, }); diff --git a/app/packages/looker/src/worker/painters/make-mask-colorer.ts b/app/packages/looker/src/worker/painters/make-mask-colorer.ts new file mode 100644 index 0000000000..c5573b37b4 --- /dev/null +++ b/app/packages/looker/src/worker/painters/make-mask-colorer.ts @@ -0,0 +1,61 @@ +import { COLOR_BY, get32BitColor } from "@fiftyone/utilities"; +import type { TypedArray } from "../../numpy"; +import { convertToHex } from "../../overlays/util"; +import type { Coloring, CustomizeColor } from "../../state"; +import { convertMaskColorsToObject } from "./utils"; + +type Colorer = { + [key in COLOR_BY]: (i: number, targets: TypedArray) => number; +}; + +const makeMaskColorer = ( + coloring: Coloring, + fieldColor: number, + setting: CustomizeColor +): Colorer => { + const cache: { [i: number]: number } = {}; + + const getColor = (i: number) => { + if (!(i in cache)) { + cache[i] = get32BitColor( + convertToHex( + coloring.targets[Math.round(Math.abs(i)) % coloring.targets.length] + ) + ); + } + + return cache[i]; + }; + + // convert the defaultMaskTargetsColors and fields maskTargetsColors into + // objects to improve performance + const defaultSetting = convertMaskColorsToObject( + coloring.defaultMaskTargetsColors + ); + const fieldSetting = convertMaskColorsToObject(setting?.maskTargetsColors); + + return { + [COLOR_BY.FIELD]: () => fieldColor, + [COLOR_BY.INSTANCE]: (i: number, targets: TypedArray) => { + return getColor(Number(targets[i])); + }, + [COLOR_BY.VALUE]: (i: number, targets: TypedArray) => { + // Attempt to find a color in the fields mask target color settings + // If not found, attempt to find a color in the default mask target + // colors + const targetString = targets[i].toString(); + const customColor = + fieldSetting?.[targetString] || defaultSetting?.[targetString]; + + // If a customized color setting is found, get the 32-bit color + // representation + if (customColor) { + return get32BitColor(convertToHex(customColor)); + } + + return getColor(targets[i]); + }, + }; +}; + +export default makeMaskColorer; diff --git a/app/packages/looker/src/worker/painters/panoptic.ts b/app/packages/looker/src/worker/painters/panoptic.ts new file mode 100644 index 0000000000..b9e6b5c663 --- /dev/null +++ b/app/packages/looker/src/worker/painters/panoptic.ts @@ -0,0 +1,84 @@ +import { COLOR_BY } from "@fiftyone/utilities"; +import type { SegmentationLabel } from "../../overlays/segmentation"; +import type { MaskTargets } from "../../state"; +import type { OverlayMask } from "../decoders/types"; +import { ARRAY_TYPES } from "../decoders/types"; +import makeMaskColorer from "./make-mask-colorer"; +import type { Painter } from "./utils"; +import { getFieldColor, getTargets } from "./utils"; + +const makeShouldClear = (maskTargets: MaskTargets) => { + const isMaskTargetsEmpty = Object.keys(maskTargets).length === 0; + return (target: number) => { + if (isMaskTargetsEmpty) { + return false; + } + + return !(target in maskTargets); + }; +}; + +const panoptic: Painter = async ({ + field, + label, + coloring, + customizeColorSetting, +}) => { + console.log(label); + if (!label?.mask) { + return; + } + + // the actual overlay that'll be painted, byte-length of width * height * 4 + // (RGBA channels) + + const maskTargets = getTargets(field, coloring); + const maskData: OverlayMask = label.mask.data; + + // target map array + const targets = new ARRAY_TYPES[maskData.arrayType](maskData.buffer); + + console.log(targets); + const setting = customizeColorSetting.find((x) => x.path === field); + const fieldColor = await getFieldColor(field, coloring, setting); + console.log(label.mask); + const overlay = new Uint32Array(label.mask.image); + let colorBy = coloring.by; + + if (maskTargets && Object.keys(maskTargets).length === 1) { + colorBy = COLOR_BY.FIELD; + } + + const shouldClear = makeShouldClear(maskTargets); + const colorer = makeMaskColorer(coloring, fieldColor, setting); + const zero = 0; + + let tick = 1; + let i = 0; + let j = 0; + + if (colorBy === COLOR_BY.INSTANCE) { + i = 1; + } + + tick = maskData.channels; + + // while loop should be fast + while (i < targets.length) { + if (targets[i] === zero) { + i += tick; + j++; + continue; + } + + if (shouldClear(Number(targets[i]))) { + targets[i] = zero; + } + + overlay[j] = colorer[colorBy](i, targets); + i += tick; + j++; + } +}; + +export default panoptic; diff --git a/app/packages/looker/src/worker/painters/segmentation.ts b/app/packages/looker/src/worker/painters/segmentation.ts index c294dad327..2f2c9b2a5c 100644 --- a/app/packages/looker/src/worker/painters/segmentation.ts +++ b/app/packages/looker/src/worker/painters/segmentation.ts @@ -1,37 +1,12 @@ import { COLOR_BY, get32BitColor, rgbToHexCached } from "@fiftyone/utilities"; -import type { OverlayMask, TypedArray } from "../../numpy"; -import { ARRAY_TYPES } from "../../numpy"; import type { SegmentationLabel } from "../../overlays/segmentation"; -import { convertToHex, isRgbMaskTargets } from "../../overlays/util"; -import type { - Coloring, - CustomizeColor, - MaskTargets, - RgbMaskTargets, -} from "../../state"; +import { isRgbMaskTargets } from "../../overlays/util"; +import type { MaskTargets, RgbMaskTargets } from "../../state"; +import type { OverlayMask, TypedArray } from "../decoders/types"; +import { ARRAY_TYPES } from "../decoders/types"; +import makeMaskColorer from "./make-mask-colorer"; import type { Painter } from "./utils"; -import { - convertMaskColorsToObject, - getRgbFromMaskData, - requestColor, -} from "./utils"; - -const usingBigInts = (targets: TypedArray) => { - return targets instanceof BigInt64Array || targets instanceof BigUint64Array; -}; - -const getTargets = (field: string, coloring: Coloring) => { - // each field may have its own target map - let maskTargets: MaskTargets = coloring.maskTargets[field]; - - // or, in the absence of field specific targets, use default mask targets - // that are dataset scoped - if (!maskTargets) { - maskTargets = coloring.defaultMaskTargets; - } - - return maskTargets; -}; +import { getFieldColor, getRgbFromMaskData, getTargets } from "./utils"; const getGuard = (maskTargets: MaskTargets) => { const isMaskTargetsEmpty = Object.keys(maskTargets).length === 0; @@ -97,7 +72,6 @@ const segmentation: Painter = async ({ // the actual overlay that'll be painted, byte-length of width * height * 4 // (RGBA channels) - const maskTargets = getTargets(field, coloring); const maskData: OverlayMask = label.mask.data; @@ -118,25 +92,11 @@ const segmentation: Painter = async ({ } const shouldClear = makeShouldClear(maskTargets); - const colorer = makeColorer(coloring, fieldColor, setting); - const zero = usingBigInts(targets) ? 0n : 0; - - let tick = 1; - let i = 0; - let j = 0; - if (label.is_panoptic) { - if (colorBy === COLOR_BY.INSTANCE) { - i = 1; - } - - tick = maskData.channels; - } - - // while loop should be fast - while (i < targets.length) { + const colorer = makeMaskColorer(coloring, fieldColor, setting); + const zero = 0; + // for loop should be fast + for (let i = 0; i < targets.length; i++) { if (targets[i] === zero) { - i += tick; - j++; continue; } @@ -144,9 +104,7 @@ const segmentation: Painter = async ({ targets[i] = zero; } - overlay[j] = colorer[colorBy](i, targets); - i += tick; - j++; + overlay[i] = colorer[colorBy](i, targets); } }; @@ -162,72 +120,4 @@ const makeShouldClear = (maskTargets: MaskTargets) => { }; }; -type Colorer = { - [key in COLOR_BY]: (i: number, targets: TypedArray) => number; -}; - -const makeColorer = ( - coloring: Coloring, - fieldColor: number, - setting: CustomizeColor -): Colorer => { - const cache: { [i: number]: number } = {}; - - const getColor = (i: number) => { - if (!(i in cache)) { - cache[i] = get32BitColor( - convertToHex( - coloring.targets[Math.round(Math.abs(i)) % coloring.targets.length] - ) - ); - } - - return cache[i]; - }; - - // convert the defaultMaskTargetsColors and fields maskTargetsColors into - // objects to improve performance - const defaultSetting = convertMaskColorsToObject( - coloring.defaultMaskTargetsColors - ); - const fieldSetting = convertMaskColorsToObject(setting?.maskTargetsColors); - - return { - [COLOR_BY.FIELD]: () => fieldColor, - [COLOR_BY.INSTANCE]: (i: number, targets: TypedArray) => { - return getColor(Number(targets[i])); - }, - [COLOR_BY.VALUE]: (i: number, targets: TypedArray) => { - // Attempt to find a color in the fields mask target color settings - // If not found, attempt to find a color in the default mask target - // colors - const targetString = targets[i].toString(); - const customColor = - fieldSetting?.[targetString] || defaultSetting?.[targetString]; - - // If a customized color setting is found, get the 32-bit color - // representation - if (customColor) { - return get32BitColor(convertToHex(customColor)); - } - - return getColor(targets[i]); - }, - }; -}; - -const getFieldColor = async ( - field: string, - coloring: Coloring, - setting: CustomizeColor -) => { - // if field color has valid custom settings, use the custom field color - // convert the color into hex code, since it could be a color name - // (e.g. yellowgreen) - const fieldColorString = setting?.fieldColor - ? setting.fieldColor - : await requestColor(coloring.pool, coloring.seed, field); - return get32BitColor(convertToHex(fieldColorString)); -}; - export default segmentation; diff --git a/app/packages/looker/src/worker/painters/utils.ts b/app/packages/looker/src/worker/painters/utils.ts index ea32e98324..d168e067a8 100644 --- a/app/packages/looker/src/worker/painters/utils.ts +++ b/app/packages/looker/src/worker/painters/utils.ts @@ -1,5 +1,5 @@ +import { get32BitColor } from "@fiftyone/utilities"; import colorString from "color-string"; -import type { TypedArray } from "../../numpy"; import type { RegularLabel } from "../../overlays/base"; import type { Coloring, @@ -7,7 +7,9 @@ import type { CustomizeColor, LabelTagColor, MaskColorInput, + MaskTargets, } from "../../state"; +import type { IntermediateMask, TypedArray } from "../decoders/types"; import type { ReaderMethod } from "../types"; export type Painter