license import * as tfliteWebAPIClient from '../tflite_web_api_client' ; import { ImageSegmenter as TaskLibraryImageSegmenter } from '../types/image_segmenter' ; import { BaseTaskLibraryClient , CommonTaskLibraryOptions , getDefaultNumThreads } from './common' ; /** Different output types. */ export enum OutputType { CATEGORY_MASK = 1.0 , CONFIDENCE_MASK = 2.0 , UNSPECIFIED = 0.0 } /** ImageSegmenter options. */ export interface ImageSegmenterOptions extends CommonTaskLibraryOptions { outputType ? : OutputType ; } /** A single segmentation in the result. */ export interface Segmentation { /**
     * The width of the mask. This is an intrinsic parameter of the model being
     * used, and does not depend on the input image dimensions.
     */ width : number ; /**
     * The height of the mask. This is an intrinsic parameter of the model being
     * used, and does not depend on the input image dimensions.
     */ height : number ; /**
     * This is a flattened 2D-array of size `width` x `height`, in row major
     * order. The value of each pixel in this mask represents the class to which
     * the pixel in the mask belongs. See `coloredLabels` for instructions on how
     * to get pixel labels and display color.
     */ categoryMask : Uint8Array ; /**
     * The list of colored labels for all the supported categories. Depending on
     * which is present, this list is in 1:1 correspondence with:
     * - * `categoryMask` pixel values, i.e. a pixel with value `i` is associated
     * with `coloredLabels[i]`,
     * - `confidenceMasks` indices, i.e. `confidence_masks[i]` is associated with
     * `coloredLabels[i]` (TODO: to be supported here).
     */ coloredLabels : ColoredLabel [ ] ; } /** A label associated with an RGB color, for display purpose. */ export interface ColoredLabel { /** The red color component for the label, in the [0, 255] range. */ r : number ; /** The green color component for the label, in the [0, 255] range. */ g : number ; /** The blue color component for the label, in the [0, 255] range. */ b : number ; /**
     * The class name, as provided in the label map packed in the TFLite Model
     * Metadata.
     */ className : string ; /**
     * The display name, as provided in the label map (if available) packed in
     * the TFLite Model Metadata.
     */ displayName : string ; } /**
 * Client for ImageSegmenter TFLite Task Library.
 *
 * It is a wrapper around the underlying javascript API to make it more
 * convenient to use. See comments in the corresponding type declaration file in
 * src/types for more info.
 */ export class ImageSegmenter extends BaseTaskLibraryClient { constructor ( protected instance ) { super ( instance ) ; } static async create ( model , options ? ) { const optionsProto = new tfliteWebAPIClient . tfweb . ImageSegmenterOptions ( ) ; if ( options ) { if ( options . outputType ) { optionsProto . setOutputType ( options . outputType ) ; } if ( options . numThreads !== undefined ) { optionsProto . setNumThreads ( options . numThreads ) ; } } if ( ! options || options . numThreads === undefined ) { optionsProto . setNumThreads ( await getDefaultNumThreads ( ) ) ; } const instance = await tfliteWebAPIClient . tfweb . ImageSegmenter . create ( model , optionsProto ) ; return new ImageSegmenter ( instance ) ; } segment ( input ) { const result = this . instance . segment ( input ) ; if ( ! result ) { return [ ] ; } const segmentations = [ ] ; if ( result . getSegmentationList ( ) . length > 0 ) { result . getSegmentationList ( ) . forEach ( seg => { const coloredLabels = seg . getColoredLabelsList ( ) . map ( label => { return { r : label . getR ( ) , g : label . getG ( ) , b : label . getB ( ) , className : label . getClassName ( ) , displayName : label . getDisplayName ( ) , } ; } ) ; segmentations . push ( { width : seg . getWidth ( ) , height : seg . getHeight ( ) , categoryMask : seg . getCategoryMask_asU8 ( ) , coloredLabels , } ) ; } ) ; } return segmentations ; } }