Abderrahim-Mouhcine/aesthetic-immich-starter
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial upload: aesthetic-immich-starter

Browse Source
This commit is contained in:
Abderrahim-Mouhcine
2026-08-16 19:59:33 +01:00
commit 8938106e99
4777 changed files with 2176899 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/.venv/Lib/site-packages/mediapipe/__init__.py
+20
View File
@@ -0,0 +1,20 @@
# Copyright 2019 - 2022 The MediaPipe Authors.
#
# 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 mediapipe.tasks.python as tasks
from mediapipe.tasks.python.vision.core.image import Image
from mediapipe.tasks.python.vision.core.image import ImageFormat
__version__ = '0.10.35'
backend/.venv/Lib/site-packages/mediapipe/modules/hand_landmark/handedness.txt
+2
View File
@@ -0,0 +1,2 @@
Left
Right
backend/.venv/Lib/site-packages/mediapipe/modules/objectron/object_detection_oidv4_labelmap.txt
+24
View File
@@ -0,0 +1,24 @@
???
Bicycle
Boot
Laptop
Person
Chair
Cattle
Desk
Cat
Computer mouse
Computer monitor
Box
Mug
Coffee cup
Stationary bicycle
Table
Bottle
High heels
Vehicle
Footwear
Dog
Book
Camera
Car
backend/.venv/Lib/site-packages/mediapipe/tasks/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""Copyright 2022 The MediaPipe Authors.
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.
"""
backend/.venv/Lib/site-packages/mediapipe/tasks/c/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Empty __init__.py file"""
backend/.venv/Lib/site-packages/mediapipe/tasks/c/libmediapipe.dll
BIN
View File
Binary file not shown.
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/image_segmenter_metadata_schema.fbs
+59
View File
@@ -0,0 +1,59 @@
// Copyright 2023 The MediaPipe Authors.
//
// 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.
namespace mediapipe.tasks;
// Image segmenter metadata contains information specific for the image
// segmentation task. The metadata can be added in
// SubGraphMetadata.custom_metadata [1] in model metadata.
// [1]: https://github.com/google/mediapipe/blob/46b5c4012d2ef76c9d92bb0d88a6b107aee83814/mediapipe/tasks/metadata/metadata_schema.fbs#L685
// ImageSegmenterOptions.min_parser_version indicates the minimum necessary
// image segmenter metadata parser version to fully understand all fields in a
// given metadata flatbuffer. This min_parser_version is specific for the
// image segmenter metadata defined in this schema file.
//
// New fields and types will have associated comments with the schema version
// for which they were added.
//
// Schema Semantic version: 1.0.0
// This indicates the flatbuffer compatibility. The number will bump up when a
// break change is applied to the schema, such as removing fields or adding new
// fields to the middle of a table.
file_identifier "V001";
// History:
// 1.0.0 - Initial version.
// Supported activation functions.
enum Activation: byte {
NONE = 0,
SIGMOID = 1,
SOFTMAX = 2
}
table ImageSegmenterOptions {
// The activation function of the output layer in the image segmenter.
activation: Activation;
// The minimum necessary image segmenter metadata parser version to fully
// understand all fields in a given metadata flatbuffer. This field is
// automatically populated by the MetadataPopulator when the metadata is
// populated into a TFLite model. This min_parser_version is specific for the
// image segmenter metadata defined in this schema file.
min_parser_version:string;
}
root_type ImageSegmenterOptions;
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/image_segmenter_metadata_schema_py_generated.py
+108
View File
@@ -0,0 +1,108 @@
import flatbuffers
# automatically generated by the FlatBuffers compiler, do not modify
# namespace: tasks
from flatbuffers.compat import import_numpy
np = import_numpy()
class Activation(object):
NONE = 0
SIGMOID = 1
SOFTMAX = 2
class ImageSegmenterOptions(object):
__slots__ = ['_tab']
@classmethod
def GetRootAs(cls, buf, offset=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
x = ImageSegmenterOptions()
x.Init(buf, n + offset)
return x
@classmethod
def GetRootAsImageSegmenterOptions(cls, buf, offset=0):
"""This method is deprecated. Please switch to GetRootAs."""
return cls.GetRootAs(buf, offset)
@classmethod
def ImageSegmenterOptionsBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x56\x30\x30\x31", size_prefixed=size_prefixed)
# ImageSegmenterOptions
def Init(self, buf, pos):
self._tab = flatbuffers.table.Table(buf, pos)
# ImageSegmenterOptions
def Activation(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int8Flags, o + self._tab.Pos)
return 0
# ImageSegmenterOptions
def MinParserVersion(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
if o != 0:
return self._tab.String(o + self._tab.Pos)
return None
def ImageSegmenterOptionsStart(builder):
builder.StartObject(2)
def ImageSegmenterOptionsAddActivation(builder, activation):
builder.PrependInt8Slot(0, activation, 0)
def ImageSegmenterOptionsAddMinParserVersion(builder, minParserVersion):
builder.PrependUOffsetTRelativeSlot(1, flatbuffers.number_types.UOffsetTFlags.py_type(minParserVersion), 0)
def ImageSegmenterOptionsEnd(builder):
return builder.EndObject()
class ImageSegmenterOptionsT(object):
# ImageSegmenterOptionsT
def __init__(self):
self.activation = 0 # type: int
self.minParserVersion = None # type: str
@classmethod
def InitFromBuf(cls, buf, pos):
imageSegmenterOptions = ImageSegmenterOptions()
imageSegmenterOptions.Init(buf, pos)
return cls.InitFromObj(imageSegmenterOptions)
@classmethod
def InitFromPackedBuf(cls, buf, pos=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
return cls.InitFromBuf(buf, pos+n)
@classmethod
def InitFromObj(cls, imageSegmenterOptions):
x = ImageSegmenterOptionsT()
x._UnPack(imageSegmenterOptions)
return x
# ImageSegmenterOptionsT
def _UnPack(self, imageSegmenterOptions):
if imageSegmenterOptions is None:
return
self.activation = imageSegmenterOptions.Activation()
self.minParserVersion = imageSegmenterOptions.MinParserVersion()
# ImageSegmenterOptionsT
def Pack(self, builder):
if self.minParserVersion is not None:
minParserVersion = builder.CreateString(self.minParserVersion)
ImageSegmenterOptionsStart(builder)
ImageSegmenterOptionsAddActivation(builder, self.activation)
if self.minParserVersion is not None:
ImageSegmenterOptionsAddMinParserVersion(builder, minParserVersion)
imageSegmenterOptions = ImageSegmenterOptionsEnd(builder)
return imageSegmenterOptions
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/metadata_schema.fbs
+732
View File
@@ -0,0 +1,732 @@
// Copyright 2022 The MediaPipe Authors.
//
// 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.
namespace tflite;
// TFLite metadata contains both human readable and machine readable information
// about what the model does and how to use the model. It can be used as a
// README file, which elaborates the details of the model, each input/ouput
// tensor, and each associated file.
//
// An important use case of TFLite metadata is the TFLite codegen tool, which
// automatically generates the model interface based on the properties of the
// model and the tensors. The model interface provides high-level APIs to
// interact with the model, such as preprocessing the input data and running
// inferences.
//
// Entries marked with "<Codegen usage>" are used in TFLite codegen tool to
// generate the model interface. It is recommended to fill in at least those
// enties to boost the codegen performance.
// The Metadata schema is versioned by the Semantic versioning number, such as
// MAJOR.MINOR.PATCH. It tracks the schema changes according to the rules below:
// * Bump up the MAJOR number when making potentially backwards incompatible
// changes. It must be incremented if the new changes break the backwards
// compatibility. It may also include minor and patch level changes as
// needed. The true backwards compatibility is indicated by the file
// identifier.
// * Bump up the MINOR number when making backwards compatible updates for
// major features, such as supporting new content types or adding new
// processing units.
// * Bump up the PATCH number when making small backwards compatible changes,
// such as adding a new fields or deprecating certain fields (not deleting
// them).
//
// ModelMetadata.min_parser_version indicates the minimum necessary metadata
// parser version to fully understand all fields in a given metadata flatbuffer.
//
// New fields and types will have associated comments with the schema version
// for which they were added.
//
// Schema Semantic version: 1.5.0
// This indicates the flatbuffer compatibility. The number will bump up when a
// break change is applied to the schema, such as removing fields or adding new
// fields to the middle of a table.
file_identifier "M001";
// History:
// 1.0.1 - Added VOCABULARY type to AssociatedFileType.
// 1.1.0 - Added BertTokenizerOptions to ProcessUnitOptions.
// Added SentencePieceTokenizerOptions to ProcessUnitOptions.
// Added input_process_units to SubGraphMetadata.
// Added output_process_units to SubGraphMetadata.
// 1.2.0 - Added input_tensor_group to SubGraphMetadata.
// Added output_tensor_group to SubGraphMetadata.
// 1.2.1 - Added RegexTokenizerOptions to ProcessUnitOptions.
// 1.3.0 - Added AudioProperties to ContentProperties.
// 1.4.0 - Added SCANN_INDEX_FILE type to AssociatedFileType.
// 1.4.1 - Added version to AssociatedFile.
// 1.5.0 - Added CustomMetadata in SubGraphMetadata.
// File extension of any written files.
file_extension "tflitemeta";
enum AssociatedFileType : byte {
UNKNOWN = 0,
// Files such as readme.txt.
DESCRIPTIONS = 1,
// Contains a list of labels (characters separated by "\n" or in lines) that
// annotate certain axis of the tensor. For example,
// the label file in image classification. Those labels annotate the
// the output tensor, such that each value in the output tensor is the
// probability of that corresponding category specified by the label. See the
// example label file used in image classification [1].
//
// <Codegen usage>:
// If an output tensor has an associated file as TENSOR_AXIS_LABELS, return
// the output as a mapping between the labels and probability in the model
// interface.
// If multiple files of the same type are present, the first one is used by
// default; additional ones are to be distinguished from one another by their
// specified locale.
//
// TODO: Add github example link.
TENSOR_AXIS_LABELS = 2,
// Contains a list of labels (characters separated by "\n" or in lines) that
// tensor values correspond to. For example, in
// the object detection model, one of the output tensors is the detected
// classes. And each value in the tensor refers to the index of label in the
// category label file. See the example label file used in object detection
// [1].
//
// <Codegen usage>:
// If an output tensor has an associated file as TENSOR_VALUE_LABELS, convert
// the tensor values into labels, and return a list of string as the output.
// If multiple files of the same type are present, the first one is used by
// default; additional ones are to be distinguished from one another by their
// specified locale.
//
// TODO: Add github example link.
TENSOR_VALUE_LABELS = 3,
// Contains sigmoid-based score calibration parameters, formatted as CSV.
// Lines contain for each index of an output tensor the scale, slope, offset
// and (optional) min_score parameters to be used for sigmoid fitting (in this
// order and in `strtof`-compatible [1] format). Scale should be a
// non-negative value.
// A line may be left empty to default calibrated scores for this index to
// default_score.
// In summary, each line should thus contain 0, 3 or 4 comma-separated values.
//
// See the example score calibration file used in image classification [2].
//
// See documentation for ScoreCalibrationOptions for details.
//
// [1]: https://en.cppreference.com/w/c/string/byte/strtof
// TODO: Add github example link.
TENSOR_AXIS_SCORE_CALIBRATION = 4,
// Contains a list of unique words (characters separated by "\n" or in lines)
// that help to convert natural language words to embedding vectors.
//
// See the example vocab file used in text classification [1].
//
// TODO: Add github example link.
// Added in: 1.0.1
VOCABULARY = 5,
// TODO: introduce the ScaNN index file with links once the code
// is released.
// Contains on-device ScaNN index file with LevelDB format.
// Added in: 1.4.0
SCANN_INDEX_FILE = 6,
}
table AssociatedFile {
// Name of this file. Need to be exact the same as the name of the actual file
// packed into the TFLite model as a zip file.
//
// <Codegen usage>:
// Locates to the actual file in the TFLite model.
name:string;
// A description of what the file is.
description:string;
// Type of the associated file. There may be special pre/post processing for
// some types. For example in image classification, a label file of the output
// will be used to convert object index into string.
//
// <Codegen usage>:
// Determines how to process the corresponding tensor.
type:AssociatedFileType;
// An optional locale for this associated file (if applicable). It is
// recommended to use an ISO 639-1 letter code (e.g. "en" for English),
// optionally completed by a two letter region code (e.g. "en-US" for US
// English and "en-CA" for Canadian English).
// Leverage this in order to specify e.g multiple label files translated in
// different languages.
locale:string;
// Version of the file specified by model creators.
// Added in: 1.4.1
version:string;
}
// The basic content type for all tensors.
//
// <Codegen usage>:
// Input feature tensors:
// 1. Generates the method to load data from a TensorBuffer.
// 2. Creates the preprocessing logic. The default processing pipeline is:
// [NormalizeOp, QuantizeOp].
// Output feature tensors:
// 1. Generates the method to return the output data to a TensorBuffer.
// 2. Creates the post-processing logic. The default processing pipeline is:
// [DeQuantizeOp].
table FeatureProperties {
}
// The type of color space of an image.
enum ColorSpaceType : byte {
UNKNOWN = 0,
RGB = 1,
GRAYSCALE = 2,
}
table ImageSize {
width:uint;
height:uint;
}
// The properties for image tensors.
//
// <Codegen usage>:
// Input image tensors:
// 1. Generates the method to load an image from a TensorImage.
// 2. Creates the preprocessing logic. The default processing pipeline is:
// [ResizeOp, NormalizeOp, QuantizeOp].
// Output image tensors:
// 1. Generates the method to return the output data to a TensorImage.
// 2. Creates the post-processing logic. The default processing pipeline is:
// [DeQuantizeOp].
table ImageProperties {
// The color space of the image.
//
// <Codegen usage>:
// Determines how to convert the color space of a given image from users.
color_space:ColorSpaceType;
// Indicates the default value of image width and height if the tensor shape
// is dynamic. For fixed-size tensor, this size will be consistent with the
// expected size.
default_size:ImageSize;
}
// The properties for tensors representing bounding boxes.
//
// <Codegen usage>:
// Input image tensors: NA.
// Output image tensors: parses the values into a data structure that represents
// bounding boxes. For example, in the generated wrapper for Android, it returns
// the output as android.graphics.Rect objects.
enum BoundingBoxType : byte {
UNKNOWN = 0,
// Represents the bounding box by using the combination of boundaries,
// {left, top, right, bottom}.
// The default order is {left, top, right, bottom}. Other orders can be
// indicated by BoundingBoxProperties.index.
BOUNDARIES = 1,
// Represents the bounding box by using the upper_left corner, width and
// height.
// The default order is {upper_left_x, upper_left_y, width, height}. Other
// orders can be indicated by BoundingBoxProperties.index.
UPPER_LEFT = 2,
// Represents the bounding box by using the center of the box, width and
// height. The default order is {center_x, center_y, width, height}. Other
// orders can be indicated by BoundingBoxProperties.index.
CENTER = 3,
}
// The properties for audio tensors.
// Added in: 1.3.0
table AudioProperties {
// The sample rate in Hz when the audio was captured.
sample_rate:uint;
// The channel count of the audio.
channels:uint;
}
enum CoordinateType : byte {
// The coordinates are float values from 0 to 1.
RATIO = 0,
// The coordinates are integers.
PIXEL = 1,
}
table BoundingBoxProperties {
// Denotes the order of the elements defined in each bounding box type. An
// empty index array represent the default order of each bounding box type.
// For example, to denote the default order of BOUNDARIES, {left, top, right,
// bottom}, the index should be {0, 1, 2, 3}. To denote the order {left,
// right, top, bottom}, the order should be {0, 2, 1, 3}.
//
// The index array can be applied to all bounding box types to adjust the
// order of their corresponding underlying elements.
//
// <Codegen usage>:
// Indicates how to parse the bounding box values.
index:[uint];
// <Codegen usage>:
// Indicates how to parse the bounding box values.
type:BoundingBoxType;
// <Codegen usage>:
// Indicates how to convert the bounding box back to the original image in
// pixels.
coordinate_type:CoordinateType;
}
union ContentProperties {
FeatureProperties,
ImageProperties,
BoundingBoxProperties,
// Added in: 1.3.0
AudioProperties,
}
table ValueRange {
min:int;
max:int;
}
table Content {
// The properties that the content may have, indicating the type of the
// Content.
//
// <Codegen usage>:
// Indicates how to process the tensor.
content_properties:ContentProperties;
// The range of dimensions that the content corresponds to. A NULL
// "range" indicates that the content uses up all dimensions,
// except the batch axis if applied.
//
// Here are all the possible situations of how a tensor is composed.
// Case 1: The tensor is a single object, such as an image.
// For example, the input of an image classifier
// (https://www.tensorflow.org/lite/models/image_classification/overview),
// a tensor of shape [1, 224, 224, 3]. Dimensions 1 to 3 correspond to the
// image. Since dimension 0 is a batch axis, which can be ignored,
// "range" can be left as NULL.
//
// Case 2: The tensor contains multiple instances of the same object.
// For example, the output tensor of detected bounding boxes of an object
// detection model
// (https://www.tensorflow.org/lite/models/object_detection/overview).
// The tensor shape is [1, 10, 4]. Here is the what the three dimensions
// represent for:
// dimension 0: the batch axis.
// dimension 1: the 10 objects detected with the highest confidence.
// dimension 2: the bounding boxes of the 10 detected objects.
// The tensor is essentially 10 bounding boxes. In this case,
// "range" should be {min=2; max=2;}.
//
// The output tensor of scores of the above object detection model has shape
// [1, 10], where
// dimension 0: the batch axis;
// dimension 1: the scores of the 10 detected objects.
// Set "range" to the number of dimensions which is {min=2; max=2;} to denote
// that every element in the tensor is an individual content object, i.e. a
// score in this example.
//
// Another example is the pose estimation model
// (https://www.tensorflow.org/lite/models/pose_estimation/overview).
// The output tensor of heatmaps is in the shape of [1, 9, 9, 17].
// Here is the what the four dimensions represent for:
// dimension 0: the batch axis.
// dimension 1/2: the heatmap image.
// dimension 3: 17 body parts of a person.
// Even though the last axis is body part, the real content of this tensor is
// the heatmap. "range" should be [min=1; max=2].
//
// Case 3: The tensor contains multiple different objects. (Not supported by
// Content at this point).
// Sometimes a tensor may contain multiple different objects, thus different
// contents. It is very common for regression models. For example, a model
// to predict the fuel efficiency
// (https://www.tensorflow.org/tutorials/keras/regression).
// The input tensor has shape [1, 9], consisting of 9 features, such as
// "Cylinders", "Displacement", "Weight", etc. In this case, dimension 1
// contains 9 different contents. However, since these sub-dimension objects
// barely need to be specifically processed, their contents are not recorded
// in the metadata. Through, the name of each dimension can be set through
// TensorMetadata.dimension_names.
//
// Note that if it is not case 3, a tensor can only have one content type.
//
// <Codegen usage>:
// Case 1: return a processed single object of certain content type.
// Case 2: return a list of processed objects of certain content type. The
// generated model interface have API to random access those objects from
// the output.
range:ValueRange;
}
// Parameters that are used when normalizing the tensor.
table NormalizationOptions{
// mean and std are normalization parameters. Tensor values are normalized
// on a per-channel basis, by the formula
// (x - mean) / std.
// If there is only one value in mean or std, we'll propagate the value to
// all channels.
//
// Quantized models share the same normalization parameters as their
// corresponding float models. For example, an image input tensor may have
// the normalization parameter of
// mean = 127.5f and std = 127.5f.
// The image value will be normalized from [0, 255] to [-1, 1].
// Then, for quantized models, the image data should be further quantized
// according to the quantization parameters. In the case of uint8, the image
// data will be scaled back to [0, 255], while for int8, the image data will
// be scaled to [-128, 127].
//
// Both the normalization parameters and quantization parameters can be
// retrieved through the metadata extractor library.
// TODO: add link for the metadata extractor library.
// Per-channel mean of the possible values used in normalization.
//
// <Codegen usage>:
// Apply normalization to input tensors accordingly.
mean:[float];
// Per-channel standard dev. of the possible values used in normalization.
//
// <Codegen usage>:
// Apply normalization to input tensors accordingly.
std:[float];
}
// The different possible score transforms to apply to uncalibrated scores
// before applying score calibration.
enum ScoreTransformationType : byte {
// Identity function: g(x) = x.
IDENTITY = 0,
// Log function: g(x) = log(x).
LOG = 1,
// Inverse logistic function: g(x) = log(x) - log(1-x).
INVERSE_LOGISTIC = 2,
}
// Options to perform score calibration on an output tensor through sigmoid
// functions. One of the main purposes of score calibration is to make scores
// across classes comparable, so that a common threshold can be used for all
// output classes. This is meant for models producing class predictions as
// output, e.g. image classification or detection models.
//
// For each index in the output tensor, this applies:
// * `f(x) = scale / (1 + e^-(slope*g(x)+offset))` if `x > min_score` or if no
// `min_score` has been specified,
// * `f(x) = default_score` otherwise or if no scale, slope and offset have been
// specified.
// Where:
// * scale, slope, offset and (optional) min_score are index-specific parameters
// * g(x) is an index-independent transform among those defined in
// ScoreTransformationType
// * default_score is an index-independent parameter.
// An AssociatedFile with type TANSOR_AXIS_SCORE_CALIBRATION specifying the
// index-specific parameters must be associated with the corresponding
// TensorMetadata for score calibration be applied.
//
// See the example score calibration file used in image classification [1].
// TODO: Add github example link.
table ScoreCalibrationOptions {
// The function to use for transforming the uncalibrated score before
// applying score calibration.
score_transformation:ScoreTransformationType;
// The default calibrated score to apply if the uncalibrated score is
// below min_score or if no parameters were specified for a given index.
default_score:float;
}
// Performs thresholding on output tensor values, in order to filter out
// low-confidence results.
table ScoreThresholdingOptions {
// The recommended global threshold below which results are considered
// low-confidence and should be filtered out.
global_score_threshold:float;
}
// Performs Bert tokenization as in tf.text.BertTokenizer
// (https://github.com/tensorflow/text/blob/3599f6fcd2b780a2dc413b90fb9315464f10b314/docs/api_docs/python/text/BertTokenizer.md)
// Added in: 1.1.0
table BertTokenizerOptions {
// The vocabulary files used in the BertTokenizer.
vocab_file:[AssociatedFile];
}
// Performs SentencePiece tokenization as in tf.text.SentencepieceTokenizer
// (https://github.com/tensorflow/text/blob/3599f6fcd2b780a2dc413b90fb9315464f10b314/docs/api_docs/python/text/SentencepieceTokenizer.md).
// Added in: 1.1.0
table SentencePieceTokenizerOptions {
// The SentencePiece model files used in the SentencePieceTokenizer.
sentencePiece_model:[AssociatedFile];
// The optional vocabulary model files used in the SentencePieceTokenizer.
vocab_file:[AssociatedFile];
}
// Splits strings by the occurrences of delim_regex_pattern and converts the
// tokens into ids. For example, given
// delim_regex_pattern: "\W+",
// string: "Words, words, words.",
// the tokens after split are: "Words", "words", "words", "".
// And then the tokens can be converted into ids according to the vocab_file.
// Added in: 1.2.1
table RegexTokenizerOptions {
delim_regex_pattern:string;
// The vocabulary files used to convert this tokens into ids.
vocab_file:[AssociatedFile];
}
// Options that are used when processing the tensor.
union ProcessUnitOptions {
NormalizationOptions,
ScoreCalibrationOptions,
ScoreThresholdingOptions,
// Added in: 1.1.0
BertTokenizerOptions,
// Added in: 1.1.0
SentencePieceTokenizerOptions,
// Added in: 1.2.1
RegexTokenizerOptions
}
// A process unit that is used to process the tensor out-of-graph.
table ProcessUnit {
options:ProcessUnitOptions;
}
// Statistics to describe a tensor.
table Stats {
// Max and min are not currently used in tflite.support codegen. They mainly
// serve as references for users to better understand the model. They can also
// be used to validate model pre/post processing results.
// If there is only one value in max or min, we'll propagate the value to
// all channels.
// Per-channel maximum value of the tensor.
max:[float];
// Per-channel minimum value of the tensor.
min:[float];
}
// Metadata of a group of tensors. It may contain several tensors that will be
// grouped together in codegen. For example, the TFLite object detection model
// example (https://www.tensorflow.org/lite/models/object_detection/overview)
// has four outputs: classes, scores, bounding boxes, and number of detections.
// If the four outputs are bundled together using TensorGroup (for example,
// named as "detection result"), the codegen tool will generate the class,
// `DetectionResult`, which contains the class, score, and bounding box. And the
// outputs of the model will be converted to a list of `DetectionResults` and
// the number of detection. Note that the number of detection is a single
// number, therefore is inappropriate for the list of `DetectionResult`.
// Added in: 1.2.0
table TensorGroup {
// Name of tensor group.
//
// <codegen usage>:
// Name of the joint class of the tensor group.
name:string;
// Names of the tensors to group together, corresponding to
// TensorMetadata.name.
//
// <codegen usage>:
// Determines which tensors will be added to this group. All tensors in the
// group should have the same number of elements specified by Content.range.
tensor_names:[string];
}
// Detailed information of an input or output tensor.
table TensorMetadata {
// Name of the tensor.
//
// <Codegen usage>:
// The name of this tensor in the generated model interface.
name:string;
// A description of the tensor.
description:string;
// A list of names of the dimensions in this tensor. The length of
// dimension_names need to match the number of dimensions in this tensor.
//
// <Codegen usage>:
// The name of each dimension in the generated model interface. See "Case 2"
// in the comments of Content.range.
dimension_names:[string];
// The content that represents this tensor.
//
// <Codegen usage>:
// Determines how to process this tensor. See each item in ContentProperties
// for the default process units that will be applied to the tensor.
content:Content;
// The process units that are used to process the tensor out-of-graph.
//
// <Codegen usage>:
// Contains the parameters of the default processing pipeline for each content
// type, such as the normalization parameters in all content types. See the
// items under ContentProperties for the details of the default processing
// pipeline.
process_units:[ProcessUnit];
// The statistics of the tensor values.
stats:Stats;
// A list of associated files of this tensor.
//
// <Codegen usage>:
// Contains processing parameters of this tensor, such as normalization.
associated_files:[AssociatedFile];
}
table CustomMetadata {
name:string;
data:[ubyte] (force_align: 16);
}
table SubGraphMetadata {
// Name of the subgraph.
//
// Note that, since TFLite only support one subgraph at this moment, the
// Codegen tool will use the name in ModelMetadata in the generated model
// interface.
name:string;
// A description explains details about what the subgraph does.
description:string;
// Metadata of all input tensors used in this subgraph. It matches exactly
// the input tensors specified by `SubGraph.inputs` in the TFLite
// schema.fbs file[2]. The number of `TensorMetadata` in the array should
// equal to the number of indices in `SubGraph.inputs`.
//
// [2]: tensorflow/lite/schema/schema.fbs
// <Codegen usage>:
// Determines how to process the inputs.
input_tensor_metadata:[TensorMetadata];
// Metadata of all output tensors used in this subgraph. It matches exactly
// the output tensors specified by `SubGraph.outputs` in the TFLite
// schema.fbs file[2]. The number of `TensorMetadata` in the array should
// equal to the number of indices in `SubGraph.outputs`.
//
// <Codegen usage>:
// Determines how to process the outputs.
output_tensor_metadata:[TensorMetadata];
// A list of associated files of this subgraph.
associated_files:[AssociatedFile];
// Input process units of the subgraph. Some models may have complex pre and
// post processing logics where the process units do not work on one tensor at
// a time, but in a similar way of a TFLite graph. For example, in the
// MobileBert model (https://www.tensorflow.org/lite/models/bert_qa/overview),
// the inputs are: ids / mask / segment ids;
// the outputs are: end logits / start logits.
// The preprocessing converts the query string and the context string to the
// model inputs, and the post-processing converts the model outputs to the
// answer string.
// Added in: 1.1.0
input_process_units:[ProcessUnit];
// Output process units of the subgraph.
// Added in: 1.1.0
output_process_units:[ProcessUnit];
// Metadata of all input tensor groups used in this subgraph.
//
// <codegen usage>:
// Bundles the corresponding elements of the underlying input tensors together
// into a class, and converts those individual tensors into a list of the
// class objects.
// Added in: 1.2.0
input_tensor_groups:[TensorGroup];
// Metadata of all output tensor groups used in this subgraph.
//
// <codegen usage>:
// Bundles the corresponding elements of the underlying output tensors
// together into a class, and converts those individual tensors into a list of
// the class objects.
// Added in: 1.2.0
output_tensor_groups:[TensorGroup];
// A list of custom metadata.
// Added in: 1.5.0.
custom_metadata:[CustomMetadata];
}
table ModelMetadata {
// Name of the model.
//
// <Codegen usage>:
// The name of the model in the generated model interface.
name:string;
// Model description in schema.
description:string;
// Version of the model that specified by model creators.
version:string;
// Noted that, the minimum required TFLite runtime version that the model is
// compatible with, has already been added as a metadata entry in tflite
// schema. We'll decide later if we want to move it here, and keep it with
// other metadata entries.
// Metadata of all the subgraphs of the model. The 0th is assumed to be the
// main subgraph.
//
// <Codegen usage>:
// Determines how to process the inputs and outputs.
subgraph_metadata:[SubGraphMetadata];
// The person who creates this model.
author:string;
// Licenses that may apply to this model.
license:string;
// A list of associated files of this model.
associated_files:[AssociatedFile];
// The minimum metadata parser version that can fully understand the fields in
// the metadata flatbuffer. The version is effectively the largest version
// number among the versions of all the fields populated and the smallest
// compatible version indicated by the file identifier.
//
// This field is automatically populated by the MetadataPopulator when
// the metadata is populated into a TFLite model.
min_parser_version:string;
}
root_type ModelMetadata;
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/metadata_schema_py_generated.py
+3251
View File
File diff suppressed because it is too large Load Diff
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/object_detector_metadata_schema.fbs
+98
View File
@@ -0,0 +1,98 @@
// Copyright 2023 The MediaPipe Authors. All Rights Reserved.
//
// 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.
namespace mediapipe.tasks;
// ObjectDetectorOptions.min_parser_version indicates the minimum necessary
// object detector metadata parser version to fully understand all fields in a
// given metadata flatbuffer. This min_parser_version is specific for the
// object detector metadata defined in this schema file.
//
// New fields and types will have associated comments with the schema version
// for which they were added.
//
// Schema Semantic version: 1.0.0
// This indicates the flatbuffer compatibility. The number will bump up when a
// break change is applied to the schema, such as removing fields or adding new
// fields to the middle of a table.
file_identifier "V001";
// History:
// 1.0.0 - Initial version.
// A fixed size anchor.
table FixedAnchor {
x_center: float;
y_center: float;
width: float;
height: float;
}
// The schema for a list of anchors with fixed size.
table FixedAnchorsSchema {
anchors: [FixedAnchor];
}
// The ssd anchors options used in the object detector.
table SsdAnchorsOptions {
fixed_anchors_schema: FixedAnchorsSchema;
}
// The options for decoding the raw model output tensors. The options are mostly
// used in TensorsToDetectionsCalculatorOptions.
table TensorsDecodingOptions {
// The number of output classes predicted by the detection model.
num_classes: int;
// The number of output boxes predicted by the detection model.
num_boxes: int;
// The number of output values per boxes predicted by the detection
// model. The values contain bounding boxes, keypoints, etc.
num_coords: int;
// The offset of keypoint coordinates in the location tensor.
keypoint_coord_offset: int;
// The number of predicted keypoints.
num_keypoints: int;
// The dimension of each keypoint, e.g. number of values predicted for each
// keypoint.
num_values_per_keypoint: int;
// Parameters for decoding SSD detection model.
x_scale: float;
y_scale: float;
w_scale: float;
h_scale: float;
// Whether to apply exponential on box size.
apply_exponential_on_box_size: bool;
// Whether to apply sigmod function on the score.
sigmoid_score: bool;
}
table ObjectDetectorOptions {
// TODO: automatically populate min parser string.
// The minimum necessary object detector metadata parser version to fully
// understand all fields in a given metadata flatbuffer. This field is
// automatically populated by the MetadataPopulator when the metadata is
// populated into a TFLite model. This min_parser_version is specific for the
// object detector metadata defined in this schema file.
min_parser_version:string;
// The options of ssd anchors configs used by the detection model.
ssd_anchors_options:SsdAnchorsOptions;
// The tensors decoding options to convert raw tensors to detection results.
tensors_decoding_options:TensorsDecodingOptions;
}
root_type ObjectDetectorOptions;
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/object_detector_metadata_schema_py_generated.py
+674
View File
@@ -0,0 +1,674 @@
import flatbuffers
# automatically generated by the FlatBuffers compiler, do not modify
# namespace: tasks
from flatbuffers.compat import import_numpy
np = import_numpy()
class FixedAnchor(object):
__slots__ = ['_tab']
@classmethod
def GetRootAs(cls, buf, offset=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
x = FixedAnchor()
x.Init(buf, n + offset)
return x
@classmethod
def GetRootAsFixedAnchor(cls, buf, offset=0):
"""This method is deprecated. Please switch to GetRootAs."""
return cls.GetRootAs(buf, offset)
@classmethod
def FixedAnchorBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x56\x30\x30\x31", size_prefixed=size_prefixed)
# FixedAnchor
def Init(self, buf, pos):
self._tab = flatbuffers.table.Table(buf, pos)
# FixedAnchor
def XCenter(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# FixedAnchor
def YCenter(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# FixedAnchor
def Width(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(8))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# FixedAnchor
def Height(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(10))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
def FixedAnchorStart(builder):
builder.StartObject(4)
def FixedAnchorAddXCenter(builder, xCenter):
builder.PrependFloat32Slot(0, xCenter, 0.0)
def FixedAnchorAddYCenter(builder, yCenter):
builder.PrependFloat32Slot(1, yCenter, 0.0)
def FixedAnchorAddWidth(builder, width):
builder.PrependFloat32Slot(2, width, 0.0)
def FixedAnchorAddHeight(builder, height):
builder.PrependFloat32Slot(3, height, 0.0)
def FixedAnchorEnd(builder):
return builder.EndObject()
class FixedAnchorT(object):
# FixedAnchorT
def __init__(self):
self.xCenter = 0.0 # type: float
self.yCenter = 0.0 # type: float
self.width = 0.0 # type: float
self.height = 0.0 # type: float
@classmethod
def InitFromBuf(cls, buf, pos):
fixedAnchor = FixedAnchor()
fixedAnchor.Init(buf, pos)
return cls.InitFromObj(fixedAnchor)
@classmethod
def InitFromPackedBuf(cls, buf, pos=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
return cls.InitFromBuf(buf, pos+n)
@classmethod
def InitFromObj(cls, fixedAnchor):
x = FixedAnchorT()
x._UnPack(fixedAnchor)
return x
# FixedAnchorT
def _UnPack(self, fixedAnchor):
if fixedAnchor is None:
return
self.xCenter = fixedAnchor.XCenter()
self.yCenter = fixedAnchor.YCenter()
self.width = fixedAnchor.Width()
self.height = fixedAnchor.Height()
# FixedAnchorT
def Pack(self, builder):
FixedAnchorStart(builder)
FixedAnchorAddXCenter(builder, self.xCenter)
FixedAnchorAddYCenter(builder, self.yCenter)
FixedAnchorAddWidth(builder, self.width)
FixedAnchorAddHeight(builder, self.height)
fixedAnchor = FixedAnchorEnd(builder)
return fixedAnchor
class FixedAnchorsSchema(object):
__slots__ = ['_tab']
@classmethod
def GetRootAs(cls, buf, offset=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
x = FixedAnchorsSchema()
x.Init(buf, n + offset)
return x
@classmethod
def GetRootAsFixedAnchorsSchema(cls, buf, offset=0):
"""This method is deprecated. Please switch to GetRootAs."""
return cls.GetRootAs(buf, offset)
@classmethod
def FixedAnchorsSchemaBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x56\x30\x30\x31", size_prefixed=size_prefixed)
# FixedAnchorsSchema
def Init(self, buf, pos):
self._tab = flatbuffers.table.Table(buf, pos)
# FixedAnchorsSchema
def Anchors(self, j):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
x = self._tab.Vector(o)
x += flatbuffers.number_types.UOffsetTFlags.py_type(j) * 4
x = self._tab.Indirect(x)
obj = FixedAnchor()
obj.Init(self._tab.Bytes, x)
return obj
return None
# FixedAnchorsSchema
def AnchorsLength(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
return self._tab.VectorLen(o)
return 0
# FixedAnchorsSchema
def AnchorsIsNone(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
return o == 0
def FixedAnchorsSchemaStart(builder):
builder.StartObject(1)
def FixedAnchorsSchemaAddAnchors(builder, anchors):
builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(anchors), 0)
def FixedAnchorsSchemaStartAnchorsVector(builder, numElems):
return builder.StartVector(4, numElems, 4)
def FixedAnchorsSchemaEnd(builder):
return builder.EndObject()
try:
from typing import List
except:
pass
class FixedAnchorsSchemaT(object):
# FixedAnchorsSchemaT
def __init__(self):
self.anchors = None # type: List[FixedAnchorT]
@classmethod
def InitFromBuf(cls, buf, pos):
fixedAnchorsSchema = FixedAnchorsSchema()
fixedAnchorsSchema.Init(buf, pos)
return cls.InitFromObj(fixedAnchorsSchema)
@classmethod
def InitFromPackedBuf(cls, buf, pos=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
return cls.InitFromBuf(buf, pos+n)
@classmethod
def InitFromObj(cls, fixedAnchorsSchema):
x = FixedAnchorsSchemaT()
x._UnPack(fixedAnchorsSchema)
return x
# FixedAnchorsSchemaT
def _UnPack(self, fixedAnchorsSchema):
if fixedAnchorsSchema is None:
return
if not fixedAnchorsSchema.AnchorsIsNone():
self.anchors = []
for i in range(fixedAnchorsSchema.AnchorsLength()):
if fixedAnchorsSchema.Anchors(i) is None:
self.anchors.append(None)
else:
fixedAnchor_ = FixedAnchorT.InitFromObj(fixedAnchorsSchema.Anchors(i))
self.anchors.append(fixedAnchor_)
# FixedAnchorsSchemaT
def Pack(self, builder):
if self.anchors is not None:
anchorslist = []
for i in range(len(self.anchors)):
anchorslist.append(self.anchors[i].Pack(builder))
FixedAnchorsSchemaStartAnchorsVector(builder, len(self.anchors))
for i in reversed(range(len(self.anchors))):
builder.PrependUOffsetTRelative(anchorslist[i])
anchors = builder.EndVector()
FixedAnchorsSchemaStart(builder)
if self.anchors is not None:
FixedAnchorsSchemaAddAnchors(builder, anchors)
fixedAnchorsSchema = FixedAnchorsSchemaEnd(builder)
return fixedAnchorsSchema
class SsdAnchorsOptions(object):
__slots__ = ['_tab']
@classmethod
def GetRootAs(cls, buf, offset=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
x = SsdAnchorsOptions()
x.Init(buf, n + offset)
return x
@classmethod
def GetRootAsSsdAnchorsOptions(cls, buf, offset=0):
"""This method is deprecated. Please switch to GetRootAs."""
return cls.GetRootAs(buf, offset)
@classmethod
def SsdAnchorsOptionsBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x56\x30\x30\x31", size_prefixed=size_prefixed)
# SsdAnchorsOptions
def Init(self, buf, pos):
self._tab = flatbuffers.table.Table(buf, pos)
# SsdAnchorsOptions
def FixedAnchorsSchema(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
x = self._tab.Indirect(o + self._tab.Pos)
obj = FixedAnchorsSchema()
obj.Init(self._tab.Bytes, x)
return obj
return None
def SsdAnchorsOptionsStart(builder):
builder.StartObject(1)
def SsdAnchorsOptionsAddFixedAnchorsSchema(builder, fixedAnchorsSchema):
builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(fixedAnchorsSchema), 0)
def SsdAnchorsOptionsEnd(builder):
return builder.EndObject()
try:
from typing import Optional
except:
pass
class SsdAnchorsOptionsT(object):
# SsdAnchorsOptionsT
def __init__(self):
self.fixedAnchorsSchema = None # type: Optional[FixedAnchorsSchemaT]
@classmethod
def InitFromBuf(cls, buf, pos):
ssdAnchorsOptions = SsdAnchorsOptions()
ssdAnchorsOptions.Init(buf, pos)
return cls.InitFromObj(ssdAnchorsOptions)
@classmethod
def InitFromPackedBuf(cls, buf, pos=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
return cls.InitFromBuf(buf, pos+n)
@classmethod
def InitFromObj(cls, ssdAnchorsOptions):
x = SsdAnchorsOptionsT()
x._UnPack(ssdAnchorsOptions)
return x
# SsdAnchorsOptionsT
def _UnPack(self, ssdAnchorsOptions):
if ssdAnchorsOptions is None:
return
if ssdAnchorsOptions.FixedAnchorsSchema() is not None:
self.fixedAnchorsSchema = FixedAnchorsSchemaT.InitFromObj(ssdAnchorsOptions.FixedAnchorsSchema())
# SsdAnchorsOptionsT
def Pack(self, builder):
if self.fixedAnchorsSchema is not None:
fixedAnchorsSchema = self.fixedAnchorsSchema.Pack(builder)
SsdAnchorsOptionsStart(builder)
if self.fixedAnchorsSchema is not None:
SsdAnchorsOptionsAddFixedAnchorsSchema(builder, fixedAnchorsSchema)
ssdAnchorsOptions = SsdAnchorsOptionsEnd(builder)
return ssdAnchorsOptions
class TensorsDecodingOptions(object):
__slots__ = ['_tab']
@classmethod
def GetRootAs(cls, buf, offset=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
x = TensorsDecodingOptions()
x.Init(buf, n + offset)
return x
@classmethod
def GetRootAsTensorsDecodingOptions(cls, buf, offset=0):
"""This method is deprecated. Please switch to GetRootAs."""
return cls.GetRootAs(buf, offset)
@classmethod
def TensorsDecodingOptionsBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x56\x30\x30\x31", size_prefixed=size_prefixed)
# TensorsDecodingOptions
def Init(self, buf, pos):
self._tab = flatbuffers.table.Table(buf, pos)
# TensorsDecodingOptions
def NumClasses(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
return 0
# TensorsDecodingOptions
def NumBoxes(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
return 0
# TensorsDecodingOptions
def NumCoords(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(8))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
return 0
# TensorsDecodingOptions
def KeypointCoordOffset(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(10))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
return 0
# TensorsDecodingOptions
def NumKeypoints(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(12))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
return 0
# TensorsDecodingOptions
def NumValuesPerKeypoint(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(14))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
return 0
# TensorsDecodingOptions
def XScale(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(16))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# TensorsDecodingOptions
def YScale(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(18))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# TensorsDecodingOptions
def WScale(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(20))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# TensorsDecodingOptions
def HScale(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(22))
if o != 0:
return self._tab.Get(flatbuffers.number_types.Float32Flags, o + self._tab.Pos)
return 0.0
# TensorsDecodingOptions
def ApplyExponentialOnBoxSize(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(24))
if o != 0:
return bool(self._tab.Get(flatbuffers.number_types.BoolFlags, o + self._tab.Pos))
return False
# TensorsDecodingOptions
def SigmoidScore(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(26))
if o != 0:
return bool(self._tab.Get(flatbuffers.number_types.BoolFlags, o + self._tab.Pos))
return False
def TensorsDecodingOptionsStart(builder):
builder.StartObject(12)
def TensorsDecodingOptionsAddNumClasses(builder, numClasses):
builder.PrependInt32Slot(0, numClasses, 0)
def TensorsDecodingOptionsAddNumBoxes(builder, numBoxes):
builder.PrependInt32Slot(1, numBoxes, 0)
def TensorsDecodingOptionsAddNumCoords(builder, numCoords):
builder.PrependInt32Slot(2, numCoords, 0)
def TensorsDecodingOptionsAddKeypointCoordOffset(builder, keypointCoordOffset):
builder.PrependInt32Slot(3, keypointCoordOffset, 0)
def TensorsDecodingOptionsAddNumKeypoints(builder, numKeypoints):
builder.PrependInt32Slot(4, numKeypoints, 0)
def TensorsDecodingOptionsAddNumValuesPerKeypoint(builder, numValuesPerKeypoint):
builder.PrependInt32Slot(5, numValuesPerKeypoint, 0)
def TensorsDecodingOptionsAddXScale(builder, xScale):
builder.PrependFloat32Slot(6, xScale, 0.0)
def TensorsDecodingOptionsAddYScale(builder, yScale):
builder.PrependFloat32Slot(7, yScale, 0.0)
def TensorsDecodingOptionsAddWScale(builder, wScale):
builder.PrependFloat32Slot(8, wScale, 0.0)
def TensorsDecodingOptionsAddHScale(builder, hScale):
builder.PrependFloat32Slot(9, hScale, 0.0)
def TensorsDecodingOptionsAddApplyExponentialOnBoxSize(builder, applyExponentialOnBoxSize):
builder.PrependBoolSlot(10, applyExponentialOnBoxSize, 0)
def TensorsDecodingOptionsAddSigmoidScore(builder, sigmoidScore):
builder.PrependBoolSlot(11, sigmoidScore, 0)
def TensorsDecodingOptionsEnd(builder):
return builder.EndObject()
class TensorsDecodingOptionsT(object):
# TensorsDecodingOptionsT
def __init__(self):
self.numClasses = 0 # type: int
self.numBoxes = 0 # type: int
self.numCoords = 0 # type: int
self.keypointCoordOffset = 0 # type: int
self.numKeypoints = 0 # type: int
self.numValuesPerKeypoint = 0 # type: int
self.xScale = 0.0 # type: float
self.yScale = 0.0 # type: float
self.wScale = 0.0 # type: float
self.hScale = 0.0 # type: float
self.applyExponentialOnBoxSize = False # type: bool
self.sigmoidScore = False # type: bool
@classmethod
def InitFromBuf(cls, buf, pos):
tensorsDecodingOptions = TensorsDecodingOptions()
tensorsDecodingOptions.Init(buf, pos)
return cls.InitFromObj(tensorsDecodingOptions)
@classmethod
def InitFromPackedBuf(cls, buf, pos=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
return cls.InitFromBuf(buf, pos+n)
@classmethod
def InitFromObj(cls, tensorsDecodingOptions):
x = TensorsDecodingOptionsT()
x._UnPack(tensorsDecodingOptions)
return x
# TensorsDecodingOptionsT
def _UnPack(self, tensorsDecodingOptions):
if tensorsDecodingOptions is None:
return
self.numClasses = tensorsDecodingOptions.NumClasses()
self.numBoxes = tensorsDecodingOptions.NumBoxes()
self.numCoords = tensorsDecodingOptions.NumCoords()
self.keypointCoordOffset = tensorsDecodingOptions.KeypointCoordOffset()
self.numKeypoints = tensorsDecodingOptions.NumKeypoints()
self.numValuesPerKeypoint = tensorsDecodingOptions.NumValuesPerKeypoint()
self.xScale = tensorsDecodingOptions.XScale()
self.yScale = tensorsDecodingOptions.YScale()
self.wScale = tensorsDecodingOptions.WScale()
self.hScale = tensorsDecodingOptions.HScale()
self.applyExponentialOnBoxSize = tensorsDecodingOptions.ApplyExponentialOnBoxSize()
self.sigmoidScore = tensorsDecodingOptions.SigmoidScore()
# TensorsDecodingOptionsT
def Pack(self, builder):
TensorsDecodingOptionsStart(builder)
TensorsDecodingOptionsAddNumClasses(builder, self.numClasses)
TensorsDecodingOptionsAddNumBoxes(builder, self.numBoxes)
TensorsDecodingOptionsAddNumCoords(builder, self.numCoords)
TensorsDecodingOptionsAddKeypointCoordOffset(builder, self.keypointCoordOffset)
TensorsDecodingOptionsAddNumKeypoints(builder, self.numKeypoints)
TensorsDecodingOptionsAddNumValuesPerKeypoint(builder, self.numValuesPerKeypoint)
TensorsDecodingOptionsAddXScale(builder, self.xScale)
TensorsDecodingOptionsAddYScale(builder, self.yScale)
TensorsDecodingOptionsAddWScale(builder, self.wScale)
TensorsDecodingOptionsAddHScale(builder, self.hScale)
TensorsDecodingOptionsAddApplyExponentialOnBoxSize(builder, self.applyExponentialOnBoxSize)
TensorsDecodingOptionsAddSigmoidScore(builder, self.sigmoidScore)
tensorsDecodingOptions = TensorsDecodingOptionsEnd(builder)
return tensorsDecodingOptions
class ObjectDetectorOptions(object):
__slots__ = ['_tab']
@classmethod
def GetRootAs(cls, buf, offset=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
x = ObjectDetectorOptions()
x.Init(buf, n + offset)
return x
@classmethod
def GetRootAsObjectDetectorOptions(cls, buf, offset=0):
"""This method is deprecated. Please switch to GetRootAs."""
return cls.GetRootAs(buf, offset)
@classmethod
def ObjectDetectorOptionsBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x56\x30\x30\x31", size_prefixed=size_prefixed)
# ObjectDetectorOptions
def Init(self, buf, pos):
self._tab = flatbuffers.table.Table(buf, pos)
# ObjectDetectorOptions
def MinParserVersion(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
if o != 0:
return self._tab.String(o + self._tab.Pos)
return None
# ObjectDetectorOptions
def SsdAnchorsOptions(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
if o != 0:
x = self._tab.Indirect(o + self._tab.Pos)
obj = SsdAnchorsOptions()
obj.Init(self._tab.Bytes, x)
return obj
return None
# ObjectDetectorOptions
def TensorsDecodingOptions(self):
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(8))
if o != 0:
x = self._tab.Indirect(o + self._tab.Pos)
obj = TensorsDecodingOptions()
obj.Init(self._tab.Bytes, x)
return obj
return None
def ObjectDetectorOptionsStart(builder):
builder.StartObject(3)
def ObjectDetectorOptionsAddMinParserVersion(builder, minParserVersion):
builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(minParserVersion), 0)
def ObjectDetectorOptionsAddSsdAnchorsOptions(builder, ssdAnchorsOptions):
builder.PrependUOffsetTRelativeSlot(1, flatbuffers.number_types.UOffsetTFlags.py_type(ssdAnchorsOptions), 0)
def ObjectDetectorOptionsAddTensorsDecodingOptions(builder, tensorsDecodingOptions):
builder.PrependUOffsetTRelativeSlot(2, flatbuffers.number_types.UOffsetTFlags.py_type(tensorsDecodingOptions), 0)
def ObjectDetectorOptionsEnd(builder):
return builder.EndObject()
try:
from typing import Optional
except:
pass
class ObjectDetectorOptionsT(object):
# ObjectDetectorOptionsT
def __init__(self):
self.minParserVersion = None # type: str
self.ssdAnchorsOptions = None # type: Optional[SsdAnchorsOptionsT]
self.tensorsDecodingOptions = None # type: Optional[TensorsDecodingOptionsT]
@classmethod
def InitFromBuf(cls, buf, pos):
objectDetectorOptions = ObjectDetectorOptions()
objectDetectorOptions.Init(buf, pos)
return cls.InitFromObj(objectDetectorOptions)
@classmethod
def InitFromPackedBuf(cls, buf, pos=0):
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, pos)
return cls.InitFromBuf(buf, pos+n)
@classmethod
def InitFromObj(cls, objectDetectorOptions):
x = ObjectDetectorOptionsT()
x._UnPack(objectDetectorOptions)
return x
# ObjectDetectorOptionsT
def _UnPack(self, objectDetectorOptions):
if objectDetectorOptions is None:
return
self.minParserVersion = objectDetectorOptions.MinParserVersion()
if objectDetectorOptions.SsdAnchorsOptions() is not None:
self.ssdAnchorsOptions = SsdAnchorsOptionsT.InitFromObj(objectDetectorOptions.SsdAnchorsOptions())
if objectDetectorOptions.TensorsDecodingOptions() is not None:
self.tensorsDecodingOptions = TensorsDecodingOptionsT.InitFromObj(objectDetectorOptions.TensorsDecodingOptions())
# ObjectDetectorOptionsT
def Pack(self, builder):
if self.minParserVersion is not None:
minParserVersion = builder.CreateString(self.minParserVersion)
if self.ssdAnchorsOptions is not None:
ssdAnchorsOptions = self.ssdAnchorsOptions.Pack(builder)
if self.tensorsDecodingOptions is not None:
tensorsDecodingOptions = self.tensorsDecodingOptions.Pack(builder)
ObjectDetectorOptionsStart(builder)
if self.minParserVersion is not None:
ObjectDetectorOptionsAddMinParserVersion(builder, minParserVersion)
if self.ssdAnchorsOptions is not None:
ObjectDetectorOptionsAddSsdAnchorsOptions(builder, ssdAnchorsOptions)
if self.tensorsDecodingOptions is not None:
ObjectDetectorOptionsAddTensorsDecodingOptions(builder, tensorsDecodingOptions)
objectDetectorOptions = ObjectDetectorOptionsEnd(builder)
return objectDetectorOptions
backend/.venv/Lib/site-packages/mediapipe/tasks/metadata/schema_py_generated.py
+18667
View File
File diff suppressed because it is too large Load Diff
backend/.venv/Lib/site-packages/mediapipe/tasks/python/__init__.py
+27
View File
@@ -0,0 +1,27 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks API."""
from . import audio
from . import components
from . import core
from . import genai
from . import text
from . import vision
BaseOptions = core.base_options.BaseOptions
# Remove unnecessary modules to avoid duplication in API docs.
del core
backend/.venv/Lib/site-packages/mediapipe/tasks/python/audio/__init__.py
+28
View File
@@ -0,0 +1,28 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks Audio API."""
import mediapipe.tasks.python.audio.core
import mediapipe.tasks.python.audio.audio_classifier
AudioClassifier = audio_classifier.AudioClassifier
AudioClassifierOptions = audio_classifier.AudioClassifierOptions
AudioClassifierResult = audio_classifier.AudioClassifierResult
RunningMode = core.audio_task_running_mode.AudioTaskRunningMode
# Remove unnecessary modules to avoid duplication in API docs.
del audio_classifier
del core
del mediapipe
backend/.venv/Lib/site-packages/mediapipe/tasks/python/audio/audio_classifier.py
+456
View File
@@ -0,0 +1,456 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe audio classifier task."""
import ctypes
import dataclasses
from typing import Callable, Optional
from mediapipe.tasks.python.audio.core import audio_record
from mediapipe.tasks.python.audio.core import audio_task_running_mode
from mediapipe.tasks.python.components.containers import audio_data
from mediapipe.tasks.python.components.containers import audio_data_c
from mediapipe.tasks.python.components.containers import classification_result
from mediapipe.tasks.python.components.containers import classification_result_c
from mediapipe.tasks.python.components.processors import classifier_options as classifier_options_lib
from mediapipe.tasks.python.components.processors import classifier_options_c
from mediapipe.tasks.python.core import async_result_dispatcher
from mediapipe.tasks.python.core import base_options as base_options_lib
from mediapipe.tasks.python.core import base_options_c
from mediapipe.tasks.python.core import mediapipe_c_bindings
from mediapipe.tasks.python.core import mediapipe_c_utils
from mediapipe.tasks.python.core import serial_dispatcher
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
AudioClassifierResult = classification_result.ClassificationResult
_AudioData = audio_data.AudioData
_BaseOptions = base_options_lib.BaseOptions
_RunningMode = audio_task_running_mode.AudioTaskRunningMode
_MICRO_SECONDS_PER_MILLISECOND = 1000
_AsyncResultDispatcher = async_result_dispatcher.AsyncResultDispatcher
_LiveStreamPacket = async_result_dispatcher.LiveStreamPacket
class AudioClassifierResultC(ctypes.Structure):
"""The C representation of a list of audio classification results."""
_fields_ = [
(
'results',
ctypes.POINTER(classification_result_c.ClassificationResultC),
),
('results_count', ctypes.c_int),
]
_C_TYPES_RESULT_CALLBACK = ctypes.CFUNCTYPE(
None, ctypes.c_int32, ctypes.POINTER(AudioClassifierResultC)
)
class AudioClassifierOptionsC(ctypes.Structure):
"""The audio classifier options used in the C API."""
_fields_ = [
('base_options', base_options_c.BaseOptionsC),
('classifier_options', classifier_options_c.ClassifierOptionsC),
('running_mode', ctypes.c_int),
('result_callback', _C_TYPES_RESULT_CALLBACK),
]
@classmethod
@doc_controls.do_not_generate_docs
def from_c_options(
cls,
base_options: base_options_c.BaseOptionsC,
classifier_options: classifier_options_c.ClassifierOptionsC,
running_mode: _RunningMode,
result_callback: _C_TYPES_RESULT_CALLBACK,
) -> 'AudioClassifierOptionsC':
"""Creates an AudioClassifierOptionsC object from the given options."""
return cls(
base_options=base_options,
classifier_options=classifier_options,
running_mode=running_mode.ctype,
result_callback=result_callback,
)
_CTYPES_SIGNATURES = (
mediapipe_c_utils.CStatusFunction(
'MpAudioClassifierCreate',
(
ctypes.POINTER(AudioClassifierOptionsC),
ctypes.POINTER(ctypes.c_void_p),
),
),
mediapipe_c_utils.CStatusFunction(
'MpAudioClassifierClassify',
(
ctypes.c_void_p,
ctypes.POINTER(audio_data_c.AudioDataC),
ctypes.POINTER(AudioClassifierResultC),
),
),
mediapipe_c_utils.CStatusFunction(
'MpAudioClassifierClassifyAsync',
(
ctypes.c_void_p,
ctypes.POINTER(audio_data_c.AudioDataC),
ctypes.c_int64,
),
),
mediapipe_c_utils.CFunction(
'MpAudioClassifierCloseResult',
[ctypes.POINTER(AudioClassifierResultC)],
None,
),
mediapipe_c_utils.CStatusFunction(
'MpAudioClassifierClose',
(ctypes.c_void_p,),
),
)
@dataclasses.dataclass
class AudioClassifierOptions:
"""Options for the audio classifier task.
Attributes:
base_options: Base options for the audio classifier task.
running_mode: The running mode of the task. Default to the audio clips mode.
Audio classifier task has two running modes: 1) The audio clips mode for
running classification on independent audio clips. 2) The audio stream
mode for running classification on the audio stream, such as from
microphone. In this mode, the "result_callback" below must be specified
to receive the classification results asynchronously.
display_names_locale: The locale to use for display names specified through
the TFLite Model Metadata.
max_results: The maximum number of top-scored classification results to
return.
score_threshold: Overrides the ones provided in the model metadata. Results
below this value are rejected.
category_allowlist: Allowlist of category names. If non-empty,
classification results whose category name is not in this set will be
filtered out. Duplicate or unknown category names are ignored. Mutually
exclusive with `category_denylist`.
category_denylist: Denylist of category names. If non-empty, classification
results whose category name is in this set will be filtered out. Duplicate
or unknown category names are ignored. Mutually exclusive with
`category_allowlist`.
result_callback: The user-defined result callback for processing audio
stream data. The result callback should only be specified when the running
mode is set to the audio stream mode.
"""
base_options: _BaseOptions
running_mode: _RunningMode = _RunningMode.AUDIO_CLIPS
display_names_locale: Optional[str] = None
max_results: Optional[int] = None
score_threshold: Optional[float] = None
category_allowlist: Optional[list[str]] = None
category_denylist: Optional[list[str]] = None
result_callback: Optional[Callable[[AudioClassifierResult, int], None]] = None
class AudioClassifier:
"""Class that performs audio classification on audio data.
This API expects a TFLite model with mandatory TFLite Model Metadata that
contains the mandatory AudioProperties of the solo input audio tensor and the
optional (but recommended) category labels as AssociatedFiles with type
TENSOR_AXIS_LABELS per output classification tensor.
Input tensor:
(kTfLiteFloat32)
- input audio buffer of size `[batch * samples]`.
- batch inference is not supported (`batch` is required to be 1).
- for multi-channel models, the channels must be interleaved.
At least one output tensor with:
(kTfLiteFloat32)
- `[1 x N]` array with `N` represents the number of categories.
- optional (but recommended) category labels as AssociatedFiles with type
TENSOR_AXIS_LABELS, containing one label per line. The first such
AssociatedFile (if any) is used to fill the `category_name` field of the
results. The `display_name` field is filled from the AssociatedFile (if
any) whose locale matches the `display_names_locale` field of the
`AudioClassifierOptions` used at creation time ("en" by default, i.e.
English). If none of these are available, only the `index` field of the
results will be filled.
"""
_lib: serial_dispatcher.SerialDispatcher
_handle: ctypes.c_void_p
_dispatcher: _AsyncResultDispatcher
_async_callback: _C_TYPES_RESULT_CALLBACK
def __init__(
self,
lib: serial_dispatcher.SerialDispatcher,
handle: ctypes.c_void_p,
dispatcher: _AsyncResultDispatcher,
async_callback: _C_TYPES_RESULT_CALLBACK,
):
"""Initializes the AudioClassifier instance.
Args:
lib: The serial dispatcher for the audio classifier task.
handle: The handle to the audio classifier task.
dispatcher: The async result dispatcher for the audio classifier task.
async_callback: The C callback for processing audio stream data.
"""
self._lib = lib
self._handle = handle
self._dispatcher = dispatcher
self._async_callback = async_callback
@classmethod
def create_from_model_path(cls, model_path: str) -> 'AudioClassifier':
"""Creates an `AudioClassifier` object from a TensorFlow Lite model and the default `AudioClassifierOptions`.
Note that the created `AudioClassifier` instance is in audio clips mode, for
classifying on independent audio clips.
Args:
model_path: Path to the model.
Returns:
`AudioClassifier` object that's created from the model file and the
default `AudioClassifierOptions`.
Raises:
ValueError: If failed to create `AudioClassifier` object from the provided
file such as invalid file path.
RuntimeError: If other types of error occurred.
"""
base_options = _BaseOptions(model_asset_path=model_path)
options = AudioClassifierOptions(
base_options=base_options, running_mode=_RunningMode.AUDIO_CLIPS
)
return cls.create_from_options(options)
@classmethod
def create_from_options(
cls, options: AudioClassifierOptions
) -> 'AudioClassifier':
"""Creates the `AudioClassifier` object from audio classifier options.
Args:
options: Options for the audio classifier task.
Returns:
`AudioClassifier` object that's created from `options`.
Raises:
ValueError: If failed to create `AudioClassifier` object from
`AudioClassifierOptions` such as missing the model.
RuntimeError: If other types of error occurred.
"""
lib = mediapipe_c_bindings.load_shared_library(_CTYPES_SIGNATURES)
def convert_result(c_result_ptr: ctypes.POINTER(AudioClassifierResultC)):
c_result = c_result_ptr[0]
if c_result.results_count == 0:
raise RuntimeError('No results returned from audio classifier.')
py_result = AudioClassifierResult.from_ctypes(c_result.results[0])
return (py_result, py_result.timestamp_ms)
dispatcher = _AsyncResultDispatcher(converter=convert_result)
c_callback = dispatcher.wrap_callback(
options.result_callback, _C_TYPES_RESULT_CALLBACK
)
ctypes_options = AudioClassifierOptionsC.from_c_options(
base_options=options.base_options.to_ctypes(),
classifier_options=classifier_options_c.convert_to_classifier_options_c(
classifier_options_lib.ClassifierOptions(
score_threshold=options.score_threshold,
category_allowlist=options.category_allowlist,
category_denylist=options.category_denylist,
display_names_locale=options.display_names_locale,
max_results=options.max_results,
)
),
running_mode=options.running_mode,
result_callback=c_callback,
)
classifier_handle_ptr = ctypes.c_void_p()
lib.MpAudioClassifierCreate(
ctypes.byref(ctypes_options), ctypes.byref(classifier_handle_ptr)
)
return AudioClassifier(
lib=lib,
handle=classifier_handle_ptr,
dispatcher=dispatcher,
async_callback=c_callback,
)
def classify(self, audio_clip: _AudioData) -> list[AudioClassifierResult]:
"""Performs audio classification on the provided audio clip.
The audio clip is represented as a MediaPipe AudioData. The method accepts
audio clips with various length and audio sample rate. It's required to
provide the corresponding audio sample rate within the `AudioData` object.
The input audio clip may be longer than what the model is able to process
in a single inference. When this occurs, the input audio clip is split into
multiple chunks starting at different timestamps. For this reason, this
function returns a vector of ClassificationResult objects, each associated
ith a timestamp corresponding to the start (in milliseconds) of the chunk
data that was classified, e.g:
ClassificationResult #0 (first chunk of data):
timestamp_ms: 0 (starts at 0ms)
classifications #0 (single head model):
category #0:
category_name: "Speech"
score: 0.6
category #1:
category_name: "Music"
score: 0.2
ClassificationResult #1 (second chunk of data):
timestamp_ms: 800 (starts at 800ms)
classifications #0 (single head model):
category #0:
category_name: "Speech"
score: 0.5
category #1:
category_name: "Silence"
score: 0.1
Args:
audio_clip: MediaPipe AudioData.
Returns:
An `AudioClassifierResult` object that contains a list of
classification result objects, each associated with a timestamp
corresponding to the start (in milliseconds) of the chunk data that was
classified.
Raises:
ValueError: If any of the input arguments is invalid, such as the sample
rate is not provided in the `AudioData` object.
RuntimeError: If audio classification failed to run.
"""
if not audio_clip.audio_format.sample_rate:
raise ValueError('Must provide the audio sample rate in audio data.')
c_result = AudioClassifierResultC()
self._lib.MpAudioClassifierClassify(
self._handle,
audio_clip.to_ctypes(),
ctypes.byref(c_result),
)
py_result = [
AudioClassifierResult.from_ctypes(c_result.results[i])
for i in range(c_result.results_count)
]
self._lib.MpAudioClassifierCloseResult(ctypes.byref(c_result))
return py_result
def classify_async(self, audio_block: _AudioData, timestamp_ms: int) -> None:
"""Sends audio data (a block in a continuous audio stream) to perform audio classification.
Only use this method when the AudioClassifier is created with the audio
stream running mode. The input timestamps should be monotonically increasing
for adjacent calls of this method. This method will return immediately after
the input audio data is accepted. The results will be available via the
`result_callback` provided in the `AudioClassifierOptions`. The
`classify_async` method is designed to process auido stream data such as
microphone input.
The input audio data may be longer than what the model is able to process
in a single inference. When this occurs, the input audio block is split
into multiple chunks. For this reason, the callback may be called multiple
times (once per chunk) for each call to this function.
The `result_callback` provides:
- An `AudioClassifierResult` object that contains a list of
classifications.
- The input timestamp in milliseconds.
Args:
audio_block: MediaPipe AudioData.
timestamp_ms: The timestamp of the input audio data in milliseconds.
Raises:
ValueError: If any of the followings:
1) The sample rate is not provided in the `AudioData` object or the
provided sample rate is inconsistent with the previously received.
2) The current input timestamp is smaller than what the audio
classifier has already processed.
"""
if not audio_block.audio_format.sample_rate:
raise ValueError('Must provide the audio sample rate in audio data.')
self._lib.MpAudioClassifierClassifyAsync(
self._handle,
audio_block.to_ctypes(),
timestamp_ms,
)
def create_audio_record(
self, num_channels: int, sample_rate: int, required_input_buffer_size: int
) -> audio_record.AudioRecord:
"""Creates an AudioRecord instance to record audio stream.
The returned AudioRecord instance is initialized and client needs to call
the appropriate method to start recording.
Note that MediaPipe Audio tasks will up/down sample automatically to fit the
sample rate required by the model. The default sample rate of the MediaPipe
pretrained audio model, Yamnet is 16kHz.
Args:
num_channels: The number of audio channels.
sample_rate: The audio sample rate.
required_input_buffer_size: The required input buffer size in number of
float elements.
Returns:
An AudioRecord instance.
Raises:
ValueError: If there's a problem creating the AudioRecord instance.
"""
return audio_record.AudioRecord(
num_channels, sample_rate, required_input_buffer_size
)
def close(self):
"""Shuts down the MediaPipe task instance."""
if self._handle:
self._lib.MpAudioClassifierClose(self._handle)
self._handle = None
self._dispatcher.close()
self._lib.close()
def __enter__(self):
"""Returns `self` upon entering the runtime context."""
return self
def __exit__(self, exc_type, exc_value, traceback):
"""Shuts down the MediaPipe task instance on exit of the context manager.
Args:
exc_type: The exception type that caused the exit.
exc_value: The exception value that caused the exit.
traceback: The exception traceback that caused the exit.
Raises:
RuntimeError: If the MediaPipe TextClassifier task failed to close.
"""
del exc_type, exc_value, traceback # Unused.
self.close()
def __del__(self):
self.close()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/audio/core/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""Copyright 2022 The MediaPipe Authors.
All Rights Reserved.
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.
"""
backend/.venv/Lib/site-packages/mediapipe/tasks/python/audio/core/audio_record.py
+125
View File
@@ -0,0 +1,125 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
"""A module to record audio in a streaming basis."""
import threading
import numpy as np
try:
import sounddevice as sd
except OSError as oe:
sd = None
sd_error = oe
except ImportError as ie:
sd = None
sd_error = ie
class AudioRecord(object):
"""A class to record audio in a streaming basis."""
def __init__(
self, channels: int, sampling_rate: int, buffer_size: int
) -> None:
"""Creates an AudioRecord instance.
Args:
channels: Number of input channels.
sampling_rate: Sampling rate in Hertz.
buffer_size: Size of the ring buffer in number of samples.
Raises:
ValueError: if any of the arguments is non-positive.
ImportError: if failed to import `sounddevice`.
OSError: if failed to load `PortAudio`.
"""
if sd is None:
raise sd_error
if channels <= 0:
raise ValueError('channels must be positive.')
if sampling_rate <= 0:
raise ValueError('sampling_rate must be positive.')
if buffer_size <= 0:
raise ValueError('buffer_size must be positive.')
self._audio_buffer = []
self._buffer_size = buffer_size
self._channels = channels
self._sampling_rate = sampling_rate
# Create a ring buffer to store the input audio.
self._buffer = np.zeros([buffer_size, channels], dtype=float)
self._lock = threading.Lock()
def audio_callback(data, *_):
"""A callback to receive recorded audio data from sounddevice."""
self._lock.acquire()
shift = len(data)
if shift > buffer_size:
self._buffer = np.copy(data[:buffer_size])
else:
self._buffer = np.roll(self._buffer, -shift, axis=0)
self._buffer[-shift:, :] = np.copy(data)
self._lock.release()
# Create an input stream to continuously capture the audio data.
self._stream = sd.InputStream(
channels=channels,
samplerate=sampling_rate,
callback=audio_callback,
)
@property
def channels(self) -> int:
return self._channels
@property
def sampling_rate(self) -> int:
return self._sampling_rate
@property
def buffer_size(self) -> int:
return self._buffer_size
def start_recording(self) -> None:
"""Starts the audio recording."""
# Clear the internal ring buffer.
self._buffer.fill(0)
# Start recording using sounddevice's InputStream.
self._stream.start()
def stop(self) -> None:
"""Stops the audio recording."""
self._stream.stop()
def read(self, size: int) -> np.ndarray:
"""Reads the latest audio data captured in the buffer.
Args:
size: Number of samples to read from the buffer.
Returns:
A NumPy array containing the audio data.
Raises:
ValueError: Raised if `size` is larger than the buffer size.
"""
if size > self._buffer_size:
raise ValueError('Cannot read more samples than the size of the buffer.')
elif size <= 0:
raise ValueError('Size must be positive.')
start_index = self._buffer_size - size
return np.copy(self._buffer[start_index:])
backend/.venv/Lib/site-packages/mediapipe/tasks/python/audio/core/audio_task_running_mode.py
+48
View File
@@ -0,0 +1,48 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""The running mode of MediaPipe Audio Tasks."""
import enum
import types
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
# The C API constants that map to the Python enum values.
_CTYPE_VALUE_MAP = types.MappingProxyType({
'AUDIO_CLIPS': 1,
'AUDIO_STREAM': 2,
})
class AudioTaskRunningMode(enum.Enum):
"""MediaPipe audio task running mode.
Attributes:
AUDIO_CLIPS: The mode for running a mediapipe audio task on independent
audio clips.
AUDIO_STREAM: The mode for running a mediapipe audio task on an audio
stream, such as from microphone.
"""
AUDIO_CLIPS = 'AUDIO_CLIPS'
AUDIO_STREAM = 'AUDIO_STREAM'
@property
@doc_controls.do_not_generate_docs
def ctype(self) -> int:
"""Generates a C API int object."""
ctype_value = _CTYPE_VALUE_MAP.get(self.value)
assert (
ctype_value is not None
), f'Unsupported audio task running mode: {self}'
return ctype_value
backend/.venv/Lib/site-packages/mediapipe/tasks/python/benchmark/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/benchmark/benchmark_utils.py
+70
View File
@@ -0,0 +1,70 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
"""Benchmark utils for MediaPipe Tasks."""
import os
import numpy as np
def nth_percentile(inference_times, percentile):
"""Calculate the nth percentile of the inference times."""
return np.percentile(inference_times, percentile)
def average(inference_times):
"""Calculate the average of the inference times."""
return np.mean(inference_times)
def get_test_data_path(test_srcdir, file_or_dirname_path: str) -> str:
"""Determine the test data path.
Args:
test_srcdir: The path to the test source directory.
file_or_dirname_path: The path to the file or directory.
Returns:
The full test data path.
"""
for directory, subdirs, files in os.walk(test_srcdir):
for f in subdirs + files:
path = os.path.join(directory, f)
if path.endswith(file_or_dirname_path):
return path
raise ValueError(
"No %s in test directory: %s." % (file_or_dirname_path, test_srcdir)
)
def get_model_path(custom_model, default_model_path):
"""Determine the model path based on the existence of the custom model.
Args:
custom_model: The path to the custom model provided by the user.
default_model_path: The path to the default model.
Returns:
The path to the model to be used.
"""
if custom_model is not None and os.path.exists(custom_model):
print(f"Using provided model: {custom_model}")
return custom_model
else:
if custom_model is not None:
print(
f"Warning: Provided model '{custom_model}' not found. "
"Using default model instead."
)
print(f"Using default model: {default_model_path}")
return default_model_path
backend/.venv/Lib/site-packages/mediapipe/tasks/python/benchmark/vision/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/benchmark/vision/benchmark.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
"""MediaPipe vision benchmarker."""
import argparse
from mediapipe.tasks.python.benchmark import benchmark_utils as bu
from mediapipe.tasks.python.benchmark.vision.core import base_vision_benchmark_api
from mediapipe.tasks.python.core import base_options
def benchmarker(benchmark_function, default_model_name):
"""Executes a benchmarking process using a specified function ann model.
Args:
benchmark_function: A callable function to be executed for benchmarking.
This function should contain the logic of the task to be benchmarked and
should be capable of utilizing a model specified by its name.
default_model_name: The name or path of the default model to be used in
the benchmarking process. This is useful when the benchmarking function
requires a model and no other model is explicitly specified.
"""
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
'--mode',
help='Benchmarking mode (e.g., "nth_percentile").',
required=False,
default='nth_percentile',
)
parser.add_argument('--model', help='Path to the model.', default=None)
parser.add_argument(
'--iterations',
help='Number of iterations for benchmarking.',
type=int,
default=100,
)
parser.add_argument(
'--percentile',
help='Percentile for benchmarking statistics.',
type=float,
default=95.0,
)
args = parser.parse_args()
# Get the model path
default_model_path = bu.get_test_data_path(
base_vision_benchmark_api.VISION_TEST_DATA_DIR, default_model_name
)
model_path = bu.get_model_path(args.model, default_model_path)
# Define a mapping of modes to their respective function argument lists
mode_args_mapping = {
'nth_percentile': {'percentile': args.percentile},
'average': {},
}
# Check if the mode is supported and get the argument dictionary
if args.mode not in mode_args_mapping:
raise ValueError(f'Unsupported benchmarking mode: {args.mode}')
mode_args = mode_args_mapping[args.mode]
# Run the benchmark for both CPU and GPU and calculate results based on mode
results = {}
for delegate_type in [
base_options.BaseOptions.Delegate.CPU,
base_options.BaseOptions.Delegate.GPU,
]:
inference_times = benchmark_function(
model_path, args.iterations, delegate_type
)
# Calculate the benchmark result based on the mode
if args.mode == 'nth_percentile':
results[delegate_type] = bu.nth_percentile(inference_times, **mode_args)
elif args.mode == 'average':
results[delegate_type] = bu.average(inference_times)
# Report benchmarking results
for delegate_type, result in results.items():
print(
f'Inference time {delegate_type} {mode_args_mapping[args.mode]}: '
f'{result:.6f} milliseconds'
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/benchmark/vision/core/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""Copyright 2023 The MediaPipe Authors.
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.
"""
backend/.venv/Lib/site-packages/mediapipe/tasks/python/benchmark/vision/core/base_vision_benchmark_api.py
+40
View File
@@ -0,0 +1,40 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
"""MediaPipe vision benchmark base api."""
import time
VISION_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
def benchmark_task(func, image, n_iterations):
"""Collect inference times for a given task after benchmarking.
Args:
func: The task function used for benchmarking.
image: The input MediaPipe Image.
n_iterations: Number of iterations to run the benchmark.
Returns:
List of inference times in milliseconds.
"""
inference_times = []
for _ in range(n_iterations):
start_time_ns = time.time_ns()
# Run the method for the task (e.g., classify)
func(image)
end_time_ns = time.time_ns()
inference_times.append((end_time_ns - start_time_ns) / 1_000_000)
return inference_times
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/__init__.py
+53
View File
@@ -0,0 +1,53 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks Components Containers API."""
import mediapipe.tasks.python.components.containers.audio_data
import mediapipe.tasks.python.components.containers.bounding_box
import mediapipe.tasks.python.components.containers.category
import mediapipe.tasks.python.components.containers.classification_result
import mediapipe.tasks.python.components.containers.detections
import mediapipe.tasks.python.components.containers.embedding_result
import mediapipe.tasks.python.components.containers.landmark
import mediapipe.tasks.python.components.containers.landmark_detection_result
import mediapipe.tasks.python.components.containers.rect
AudioDataFormat = audio_data.AudioDataFormat
AudioData = audio_data.AudioData
BoundingBox = bounding_box.BoundingBox
Category = category.Category
Classifications = classification_result.Classifications
ClassificationResult = classification_result.ClassificationResult
Detection = detections.Detection
DetectionResult = detections.DetectionResult
Embedding = embedding_result.Embedding
EmbeddingResult = embedding_result.EmbeddingResult
Landmark = landmark.Landmark
NormalizedLandmark = landmark.NormalizedLandmark
LandmarksDetectionResult = landmark_detection_result.LandmarksDetectionResult
Rect = rect.Rect
NormalizedRect = rect.NormalizedRect
# Remove unnecessary modules to avoid duplication in API docs.
del audio_data
del bounding_box
del category
del classification_result
del detections
del embedding_result
del landmark
del landmark_detection_result
del rect
del mediapipe
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/audio_data.py
+152
View File
@@ -0,0 +1,152 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe audio data."""
import ctypes
import dataclasses
from typing import Optional
import numpy as np
from mediapipe.tasks.python.components.containers import audio_data_c
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class AudioDataFormat:
"""Audio format metadata.
Attributes:
num_channels: the number of channels of the audio data.
sample_rate: the audio sample rate.
"""
num_channels: int = 1
sample_rate: Optional[float] = None
class AudioData(object):
"""MediaPipe Tasks' audio container."""
def __init__(
self,
buffer_length: int,
audio_format: AudioDataFormat = AudioDataFormat()
) -> None:
"""Initializes the `AudioData` object.
Args:
buffer_length: the length of the audio buffer.
audio_format: the audio format metadata.
"""
self._audio_format = audio_format
self._buffer = np.zeros([buffer_length, self._audio_format.num_channels],
dtype=np.float32)
def clear(self):
"""Clears the internal buffer and fill it with zeros."""
self._buffer.fill(0)
def load_from_array(self,
src: np.ndarray,
offset: int = 0,
size: int = -1) -> None:
"""Loads the audio data from a NumPy array.
Args:
src: A NumPy source array contains the input audio.
offset: An optional offset for loading a slice of the `src` array to the
buffer.
size: An optional size parameter denoting the number of samples to load
from the `src` array.
Raises:
ValueError: If the input array has an incorrect shape or if
`offset` + `size` exceeds the length of the `src` array.
"""
if len(src.shape) == 1:
if self._audio_format.num_channels != 1:
raise ValueError(f"Input audio is mono, but the audio data is expected "
f"to have {self._audio_format.num_channels} channels.")
elif src.shape[1] != self._audio_format.num_channels:
raise ValueError(f"Input audio contains an invalid number of channels. "
f"Expect {self._audio_format.num_channels}.")
if size < 0:
size = len(src)
if offset + size > len(src):
raise ValueError(
f"Index out of range. offset {offset} + size {size} should be <= "
f"src's length: {len(src)}")
if len(src) >= len(self._buffer):
# If the internal buffer is shorter than the load target (src), copy
# values from the end of the src array to the internal buffer.
new_offset = offset + size - len(self._buffer)
new_size = len(self._buffer)
self._buffer = src[new_offset:new_offset + new_size].copy()
else:
# Shift the internal buffer backward and add the incoming data to the end
# of the buffer.
shift = size
self._buffer = np.roll(self._buffer, -shift, axis=0)
self._buffer[-shift:, :] = src[offset:offset + size].copy()
@classmethod
def create_from_array(cls,
src: np.ndarray,
sample_rate: Optional[float] = None) -> "AudioData":
"""Creates an `AudioData` object from a NumPy array.
Args:
src: A NumPy source array contains the input audio.
sample_rate: the optional audio sample rate.
Returns:
An `AudioData` object that contains a copy of the NumPy source array as
the data.
"""
src = src.astype(np.float32)
obj = cls(
buffer_length=src.shape[0],
audio_format=AudioDataFormat(
num_channels=1 if len(src.shape) == 1 else src.shape[1],
sample_rate=sample_rate))
obj.load_from_array(src)
return obj
@property
def audio_format(self) -> AudioDataFormat:
"""Gets the audio format of the audio."""
return self._audio_format
@property
def buffer_length(self) -> int:
"""Gets the sample count of the audio."""
return self._buffer.shape[0]
@property
def buffer(self) -> np.ndarray:
"""Gets the internal buffer."""
return self._buffer
@doc_controls.do_not_generate_docs
def to_ctypes(self) -> audio_data_c.AudioDataC:
"""Converts the object to a ctypes audio data object."""
return audio_data_c.AudioDataC(
num_channels=self.audio_format.num_channels,
sample_rate=self.audio_format.sample_rate,
audio_data=self._buffer.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
audio_data_size=self._buffer.size,
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/audio_data_c.py
+27
View File
@@ -0,0 +1,27 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""MediaPipe audio data ctypes."""
import ctypes
class AudioDataC(ctypes.Structure):
"""The audio data used in the C API."""
_fields_ = [
('num_channels', ctypes.c_int),
('sample_rate', ctypes.c_double),
('audio_data', ctypes.POINTER(ctypes.c_float)),
('audio_data_size', ctypes.c_size_t),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/bounding_box.py
+66
View File
@@ -0,0 +1,66 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Bounding box data class."""
import dataclasses
from typing import Any
from mediapipe.tasks.python.components.containers import rect_c as rect_c_lib
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class BoundingBox:
"""An integer bounding box, axis aligned.
Attributes:
origin_x: The X coordinate of the top-left corner, in pixels.
origin_y: The Y coordinate of the top-left corner, in pixels.
width: The width of the bounding box, in pixels.
height: The height of the bounding box, in pixels.
"""
origin_x: int
origin_y: int
width: int
height: int
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(cls, c_obj: rect_c_lib.RectC) -> 'BoundingBox':
"""Creates a `BoundingBox` object from a ctypes RectC struct."""
return BoundingBox(
origin_x=c_obj.left,
origin_y=c_obj.top,
width=c_obj.right - c_obj.left,
height=c_obj.bottom - c_obj.top,
)
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, BoundingBox):
return False
return (
self.origin_x == other.origin_x
and self.origin_y == other.origin_y
and self.width == other.width
and self.height == other.height
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/category.py
+97
View File
@@ -0,0 +1,97 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Category data class."""
import dataclasses
from typing import Any, Optional
from mediapipe.tasks.python.components.containers import category_c as category_c_module
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class Category:
"""A classification category.
Category is a util class, contains a label, its display name, a float
value as score, and the index of the label in the corresponding label file.
Typically it's used as the result of classification tasks.
Attributes:
index: The index of the label in the corresponding label file.
score: The probability score of this label category.
display_name: The display name of the label, which may be translated for
different locales. For example, a label, "apple", may be translated into
Spanish for display purpose, so that the `display_name` is "manzana".
category_name: The label of this category object.
"""
index: Optional[int] = None
score: Optional[float] = None
display_name: Optional[str] = None
category_name: Optional[str] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(cls, c_obj: category_c_module.CategoryC) -> "Category":
"""Creates a `Category` object from the given `CategoryC` object.
This function converts the `CategoryC` index of -1 to a Python value of None
to retain the same semantic meaning. All other values are converted as-is.
Args:
c_obj: The `CategoryC` object to be converted.
Returns:
A `Category` object.
"""
return Category(
index=c_obj.index if c_obj.index != -1 else None,
score=c_obj.score,
category_name=(
c_obj.category_name.decode("utf-8") if c_obj.category_name else None
),
display_name=(
c_obj.display_name.decode("utf-8") if c_obj.display_name else None
),
)
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, Category):
return False
return (
self.index == other.index
and self.score == other.score
and self.display_name == other.display_name
and self.category_name == other.category_name
)
@doc_controls.do_not_generate_docs
def create_list_of_categories_from_ctypes(
c_obj: category_c_module.CategoriesC,
) -> list[Category]:
"""Creates a list of `Category` objects from a `CategoriesC` object."""
return [
Category.from_ctypes(c_obj.categories[i])
for i in range(c_obj.categories_count)
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/category_c.py
+37
View File
@@ -0,0 +1,37 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""C types for Category."""
import ctypes
class CategoryC(ctypes.Structure):
"""The ctypes struct for Category."""
_fields_ = [
("index", ctypes.c_int),
("score", ctypes.c_float),
("category_name", ctypes.c_char_p),
("display_name", ctypes.c_char_p),
]
class CategoriesC(ctypes.Structure):
"""C types for a list of Categories."""
_fields_ = [
("categories", ctypes.POINTER(CategoryC)),
("categories_count", ctypes.c_uint32),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/classification_result.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright 2022 The TensorFlow Authors.
#
# 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.
"""Classifications data class."""
import dataclasses
from typing import List, Optional
from mediapipe.tasks.python.components.containers import category as category_lib
from mediapipe.tasks.python.components.containers import classification_result_c
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class Classifications:
"""Represents the classification results for a given classifier head.
Attributes:
categories: The array of predicted categories, usually sorted by descending
scores (e.g. from high to low probability).
head_index: The index of the classifier head these categories refer to. This
is useful for multi-head models.
head_name: The name of the classifier head, which is the corresponding
tensor metadata name.
"""
categories: List[category_lib.Category]
head_index: int
head_name: Optional[str] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(
cls, c_struct: classification_result_c.ClassificationsC
) -> 'Classifications':
"""Creates a `Classifications` object from the given ctypes struct."""
if c_struct.categories and c_struct.categories_count > 0:
categories = [
category_lib.Category.from_ctypes(c_struct.categories[i])
for i in range(c_struct.categories_count)
]
else:
categories = []
return Classifications(
categories=categories,
head_index=c_struct.head_index,
head_name=(
c_struct.head_name.decode('utf-8') if c_struct.head_name else None
),
)
@dataclasses.dataclass
class ClassificationResult:
"""Contains the classification results of a model.
Attributes:
classifications: A list of `Classifications` objects, each for a head of the
model.
timestamp_ms: The optional timestamp (in milliseconds) of the start of the
chunk of data corresponding to these results. This is only used for
classification on time series (e.g. audio classification). In these use
cases, the amount of data to process might exceed the maximum size that
the model can process: to solve this, the input data is split into
multiple chunks starting at different timestamps.
"""
classifications: List[Classifications]
timestamp_ms: Optional[int] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(
cls, c_struct: classification_result_c.ClassificationResultC
) -> 'ClassificationResult':
"""Creates a `ClassificationResult` object from a ctypes struct."""
if c_struct.classifications and c_struct.classifications_count > 0:
classifications = [
Classifications.from_ctypes(c_struct.classifications[i])
for i in range(c_struct.classifications_count)
]
else:
classifications = []
timestamp_ms = c_struct.timestamp_ms if c_struct.has_timestamp_ms else None
return ClassificationResult(
classifications=classifications, timestamp_ms=timestamp_ms
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/classification_result_c.py
+37
View File
@@ -0,0 +1,37 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""C types for ClassificationResult."""
import ctypes
from mediapipe.tasks.python.components.containers import category_c
class ClassificationsC(ctypes.Structure):
_fields_ = [
("categories", ctypes.POINTER(category_c.CategoryC)),
("categories_count", ctypes.c_uint32),
("head_index", ctypes.c_int),
("head_name", ctypes.c_char_p),
]
class ClassificationResultC(ctypes.Structure):
_fields_ = [
("classifications", ctypes.POINTER(ClassificationsC)),
("classifications_count", ctypes.c_uint32),
("timestamp_ms", ctypes.c_int64),
("has_timestamp_ms", ctypes.c_bool),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/detections.py
+122
View File
@@ -0,0 +1,122 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Detections data class."""
import dataclasses
from typing import Any, List, Optional
from mediapipe.tasks.python.components.containers import bounding_box as bounding_box_lib
from mediapipe.tasks.python.components.containers import category as category_lib
from mediapipe.tasks.python.components.containers import category_c as category_c_lib
from mediapipe.tasks.python.components.containers import detections_c as detections_c_lib
from mediapipe.tasks.python.components.containers import keypoint as keypoint_lib
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class Detection:
"""Represents one detected object in the object detector's results.
Attributes:
bounding_box: A BoundingBox object.
categories: A list of Category objects.
keypoints: A list of NormalizedKeypoint objects.
"""
bounding_box: bounding_box_lib.BoundingBox
categories: List[category_lib.Category]
keypoints: Optional[List[keypoint_lib.NormalizedKeypoint]] = None
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, Detection):
return False
return (
self.bounding_box == other.bounding_box
and self.categories == other.categories
and self.keypoints == other.keypoints
)
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(cls, c_obj: detections_c_lib.DetectionC) -> 'Detection':
"""Creates a `Detection` object from the given `DetectionC` object."""
c_categories = category_c_lib.CategoriesC(
categories=c_obj.categories, categories_count=c_obj.categories_count
)
py_categories = category_lib.create_list_of_categories_from_ctypes(
c_categories
)
py_bounding_box = bounding_box_lib.BoundingBox.from_ctypes(
c_obj.bounding_box
)
if c_obj.keypoints and c_obj.keypoints_count > 0:
py_keypoints = [
keypoint_lib.NormalizedKeypoint.from_ctypes(c_obj.keypoints[i])
for i in range(c_obj.keypoints_count)
]
else:
py_keypoints = None
return Detection(
bounding_box=py_bounding_box,
categories=py_categories,
keypoints=py_keypoints,
)
@dataclasses.dataclass
class DetectionResult:
"""Represents the list of detected objects.
Attributes:
detections: A list of `Detection` objects.
"""
detections: List[Detection]
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, DetectionResult):
return False
return self.detections == other.detections
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(
cls, c_obj: detections_c_lib.DetectionResultC
) -> 'DetectionResult':
"""Creates a `DetectionResult` object from a `DetectionResultC` object."""
return DetectionResult(
detections=[
Detection.from_ctypes(c_obj.detections[i])
for i in range(c_obj.detections_count)
]
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/detections_c.py
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""MediaPipe Detection Result C API types."""
import ctypes
from mediapipe.tasks.python.components.containers import category_c
from mediapipe.tasks.python.components.containers import keypoint_c
from mediapipe.tasks.python.components.containers import rect_c
class DetectionC(ctypes.Structure):
"""CTypes for a single detection."""
_fields_ = [
('categories', ctypes.POINTER(category_c.CategoryC)),
('categories_count', ctypes.c_uint32),
('bounding_box', rect_c.RectC),
('keypoints', ctypes.POINTER(keypoint_c.NormalizedKeypointC)),
('keypoints_count', ctypes.c_uint32),
]
class DetectionResultC(ctypes.Structure):
"""CTypes for the detection result."""
_fields_ = [
('detections', ctypes.POINTER(DetectionC)),
('detections_count', ctypes.c_uint32),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/embedding_result.py
+93
View File
@@ -0,0 +1,93 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Embeddings data class."""
import dataclasses
from typing import List, Optional
import numpy as np
from mediapipe.tasks.python.components.containers import embedding_result_c
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class Embedding:
"""Embedding result for a given embedder head.
Attributes:
embedding: The actual embedding, either floating-point or scalar-quantized.
head_index: The index of the embedder head that produced this embedding.
This is useful for multi-head models.
head_name: The name of the embedder head, which is the corresponding tensor
metadata name (if any). This is useful for multi-head models.
"""
embedding: np.ndarray
head_index: Optional[int] = None
head_name: Optional[str] = None
@dataclasses.dataclass
class EmbeddingResult:
"""Embedding results for a given embedder model.
Attributes:
embeddings: A list of `Embedding` objects.
timestamp_ms: The optional timestamp (in milliseconds) of the start of the
chunk of data corresponding to these results. This is only used for
embedding extraction on time series (e.g. audio embedding). In these use
cases, the amount of data to process might exceed the maximum size that
the model can process: to solve this, the input data is split into
multiple chunks starting at different timestamps.
"""
embeddings: List[Embedding]
timestamp_ms: Optional[int] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(
cls, c_result: embedding_result_c.EmbeddingResultC
) -> 'EmbeddingResult':
"""Converts a C EmbeddingResult to a Python EmbeddingResult."""
python_embeddings = []
for i in range(c_result.embeddings_count):
c_embedding = c_result.embeddings[i]
embedding = None
if c_embedding.float_embedding:
embedding_array = list(
c_embedding.float_embedding[: c_embedding.values_count]
)
embedding = np.array(embedding_array, dtype=float)
if c_embedding.quantized_embedding:
embedding_array = list(
c_embedding.quantized_embedding[: c_embedding.values_count]
)
embedding = np.array(embedding_array, dtype=np.uint8)
python_embedding = Embedding(
embedding=embedding,
head_index=c_embedding.head_index,
head_name=(
c_embedding.head_name.decode('utf-8')
if c_embedding.head_name
else None
),
)
python_embeddings.append(python_embedding)
timestamp_ms = c_result.timestamp_ms if c_result.has_timestamp_ms else None
return EmbeddingResult(
embeddings=python_embeddings, timestamp_ms=timestamp_ms
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/embedding_result_c.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Embeddings C API types."""
import ctypes
class EmbeddingC(ctypes.Structure):
"""C struct for an embedding."""
_fields_ = [
('float_embedding', ctypes.POINTER(ctypes.c_float)),
('quantized_embedding', ctypes.POINTER(ctypes.c_uint8)),
('values_count', ctypes.c_uint32),
('head_index', ctypes.c_int),
('head_name', ctypes.c_char_p),
]
class EmbeddingResultC(ctypes.Structure):
"""C struct for embedding result."""
_fields_ = [
('embeddings', ctypes.POINTER(EmbeddingC)),
('embeddings_count', ctypes.c_uint32),
('has_timestamp_ms', ctypes.c_bool),
('timestamp_ms', ctypes.c_int64),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/keypoint.py
+72
View File
@@ -0,0 +1,72 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
"""Keypoint data class."""
import dataclasses
from typing import Any, Optional
from mediapipe.tasks.python.components.containers import keypoint_c as keypoint_c_lib
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class NormalizedKeypoint:
"""A normalized keypoint.
Normalized keypoint represents a point in 2D space with x, y coordinates.
x and y are normalized to [0.0, 1.0] by the image width and height
respectively.
Attributes:
x: The x coordinates of the normalized keypoint.
y: The y coordinates of the normalized keypoint.
label: The optional label of the keypoint.
score: The score of the keypoint.
"""
x: Optional[float] = None
y: Optional[float] = None
label: Optional[str] = None
score: Optional[float] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(
cls, c_obj: keypoint_c_lib.NormalizedKeypointC
) -> 'NormalizedKeypoint':
"""Creates a `NormalizedKeypoint` object from a `NormalizedKeypointC."""
return NormalizedKeypoint(
x=c_obj.x,
y=c_obj.y,
label=c_obj.label.decode('utf-8') if c_obj.label else None,
score=c_obj.score,
)
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, NormalizedKeypoint):
return False
return (
self.x == other.x
and self.y == other.y
and self.label == other.label
and self.score == other.score
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/keypoint_c.py
+27
View File
@@ -0,0 +1,27 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""MediaPipe Keypoint C API types."""
import ctypes
class NormalizedKeypointC(ctypes.Structure):
"""A keypoint in normalized coordinates."""
_fields_ = [
('x', ctypes.c_float),
('y', ctypes.c_float),
('label', ctypes.c_char_p),
('score', ctypes.c_float),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/landmark.py
+108
View File
@@ -0,0 +1,108 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Landmark data class."""
import dataclasses
from typing import Optional
from mediapipe.tasks.python.components.containers import landmark_c as landmark_c_module
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class Landmark:
"""A landmark that can have 1 to 3 dimensions.
Use x for 1D points, (x, y) for 2D points and (x, y, z) for 3D points.
Attributes:
x: The x coordinate.
y: The y coordinate.
z: The z coordinate.
visibility: Landmark visibility. Should stay unset if not supported. Float
score of whether landmark is visible or occluded by other objects.
Landmark considered as invisible also if it is not present on the screen
(out of scene bounds). Depending on the model, visibility value is either
a sigmoid or an argument of sigmoid.
presence: Landmark presence. Should stay unset if not supported. Float score
of whether landmark is present on the scene (located within scene bounds).
Depending on the model, presence value is either a result of sigmoid or an
argument of sigmoid function to get landmark presence probability.
name: The name of the landmark.
"""
x: Optional[float] = None
y: Optional[float] = None
z: Optional[float] = None
visibility: Optional[float] = None
presence: Optional[float] = None
name: Optional[str] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(cls, c_struct: landmark_c_module.LandmarkC) -> 'Landmark':
"""Creates a `Landmark` object from the given ctypes struct."""
return Landmark(
x=c_struct.x,
y=c_struct.y,
z=c_struct.z,
visibility=c_struct.visibility if c_struct.has_visibility else None,
presence=c_struct.presence if c_struct.has_presence else None,
name=c_struct.name.decode('utf-8') if c_struct.name else None,
)
@dataclasses.dataclass
class NormalizedLandmark:
"""A normalized version of above Landmark proto.
All coordinates should be within [0, 1].
Attributes:
x: The normalized x coordinate.
y: The normalized y coordinate.
z: The normalized z coordinate.
visibility: Landmark visibility. Should stay unset if not supported. Float
score of whether landmark is visible or occluded by other objects.
Landmark considered as invisible also if it is not present on the screen
(out of scene bounds). Depending on the model, visibility value is either
a sigmoid or an argument of sigmoid.
presence: Landmark presence. Should stay unset if not supported. Float score
of whether landmark is present on the scene (located within scene bounds).
Depending on the model, presence value is either a result of sigmoid or an
argument of sigmoid function to get landmark presence probability.
name: The name of the landmark.
"""
x: Optional[float] = None
y: Optional[float] = None
z: Optional[float] = None
visibility: Optional[float] = None
presence: Optional[float] = None
name: Optional[str] = None
@classmethod
@doc_controls.do_not_generate_docs
def from_ctypes(
cls, c_struct: landmark_c_module.NormalizedLandmarkC
) -> 'NormalizedLandmark':
"""Creates a `NormalizedLandmark` object from the given ctypes struct."""
return NormalizedLandmark(
x=c_struct.x,
y=c_struct.y,
z=c_struct.z,
visibility=c_struct.visibility if c_struct.has_visibility else None,
presence=c_struct.presence if c_struct.has_presence else None,
name=c_struct.name.decode('utf-8') if c_struct.name else None,
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/landmark_c.py
+65
View File
@@ -0,0 +1,65 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""C types for Landmark."""
import ctypes
class LandmarkC(ctypes.Structure):
"""The ctypes struct for Landmark."""
_fields_ = [
('x', ctypes.c_float),
('y', ctypes.c_float),
('z', ctypes.c_float),
('has_visibility', ctypes.c_bool),
('visibility', ctypes.c_float),
('has_presence', ctypes.c_bool),
('presence', ctypes.c_float),
('name', ctypes.c_char_p),
]
class LandmarksC(ctypes.Structure):
"""The ctypes struct for a list of Landmarks."""
_fields_ = [
('landmarks', ctypes.POINTER(LandmarkC)),
('landmarks_count', ctypes.c_uint32),
]
class NormalizedLandmarkC(ctypes.Structure):
"""The ctypes struct for NormalizedLandmark."""
_fields_ = [
('x', ctypes.c_float),
('y', ctypes.c_float),
('z', ctypes.c_float),
('has_visibility', ctypes.c_bool),
('visibility', ctypes.c_float),
('has_presence', ctypes.c_bool),
('presence', ctypes.c_float),
('name', ctypes.c_char_p),
]
class NormalizedLandmarksC(ctypes.Structure):
"""The ctypes struct for a list of NormalizedLandmarks."""
_fields_ = [
('landmarks', ctypes.POINTER(NormalizedLandmarkC)),
('landmarks_count', ctypes.c_uint32),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/landmark_detection_result.py
+42
View File
@@ -0,0 +1,42 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Landmarks Detection Result data class."""
import dataclasses
from typing import Optional, List
from mediapipe.tasks.python.components.containers import category as category_module
from mediapipe.tasks.python.components.containers import landmark as landmark_module
from mediapipe.tasks.python.components.containers import rect as rect_module
_Category = category_module.Category
_NormalizedLandmark = landmark_module.NormalizedLandmark
_Landmark = landmark_module.Landmark
@dataclasses.dataclass
class LandmarksDetectionResult:
"""Represents the landmarks detection result.
Attributes:
landmarks: A list of `NormalizedLandmark` objects.
categories: A list of `Category` objects.
world_landmarks: A list of `Landmark` objects.
rect: A `NormalizedRect` object.
"""
landmarks: Optional[List[_NormalizedLandmark]]
categories: Optional[List[_Category]]
world_landmarks: Optional[List[_Landmark]]
rect: rect_module.NormalizedRect
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/matrix_c.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""C types for Matrix."""
import ctypes
import numpy as np
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
class MatrixC(ctypes.Structure):
"""The ctypes struct for Matrix."""
_fields_ = [
('rows', ctypes.c_uint32),
('cols', ctypes.c_uint32),
('data', ctypes.POINTER(ctypes.c_float)),
]
@doc_controls.do_not_generate_docs
def to_numpy(self) -> np.ndarray:
"""Copies the MatrixC struct to a numpy array.
The data is reshaped from column-major order to (cols, rows) and then
transposed to a row-major numpy array. The underlying data is copied to
ensure the data is not freed when the underlying C Matrix is closed.
Returns:
A numpy array representing the MatrixC struct.
"""
if not self.data or self.rows == 0 or self.cols == 0:
return np.empty((0,))
np_array = np.ctypeslib.as_array(
self.data, shape=(self.cols, self.rows)
).copy()
return np_array.T
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/rect.py
+123
View File
@@ -0,0 +1,123 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Rect data class."""
import dataclasses
from typing import Any, Optional
from mediapipe.tasks.python.components.containers import rect_c as rect_c_module
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
@dataclasses.dataclass
class Rect:
"""A rectangle, used as part of detection results.
Absolute coordinates. The origin is on the top-left corner of the image.
Attributes:
left: The X coordinate of the left side of the rectangle.
top: The Y coordinate of the top of the rectangle.
bottom: The Y coordinate of the bottom of the rectangle.
right: The X coordinate of the right side of the rectangle.
"""
left: int
top: int
bottom: int
right: int
@doc_controls.do_not_generate_docs
def to_ctypes(self) -> rect_c_module.RectC:
"""Generates a C API RectC object."""
return rect_c_module.RectC(
left=self.left, top=self.top, bottom=self.bottom, right=self.right
)
@dataclasses.dataclass
class RectF:
"""A rectangle, used as part as input region-of-interest.
The coordinates are normalized wrt the image dimensions, i.e. generally in
[0,1] but they may exceed these bounds if describing a region overlapping the
image. The origin is on the top-left corner of the image.
Attributes:
left: The X coordinate of the left side of the rectangle.
top: The Y coordinate of the top of the rectangle.
bottom: The Y coordinate of the bottom of the rectangle.
right: The X coordinate of the right side of the rectangle.
"""
left: float
top: float
bottom: float
right: float
@doc_controls.do_not_generate_docs
def to_ctypes(self) -> rect_c_module.RectFC:
"""Generates a C API RectFC object."""
return rect_c_module.RectFC(
left=self.left, top=self.top, bottom=self.bottom, right=self.right
)
@dataclasses.dataclass
class NormalizedRect:
"""A rectangle with rotation in normalized coordinates.
Location of the center of the rectangle in image coordinates. The (0.0, 0.0)
point is at the (top, left) corner.
The values of box center location and size are within [0, 1].
Attributes:
x_center: The normalized X coordinate of the rectangle, in image
coordinates.
y_center: The normalized Y coordinate of the rectangle, in image
coordinates.
width: The width of the rectangle.
height: The height of the rectangle.
rotation: Rotation angle is clockwise in radians.
rect_id: Optional unique id to help associate different rectangles to each
other.
"""
x_center: float
y_center: float
width: float
height: float
rotation: Optional[float] = 0.0
rect_id: Optional[int] = None
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, NormalizedRect):
return False
return (
self.x_center == other.x_center
and self.y_center == other.y_center
and self.width == other.width
and self.height == other.height
and self.rotation == other.rotation
and self.rect_id == other.rect_id
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/containers/rect_c.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""MediaPipe Rect C API types."""
import ctypes
class RectC(ctypes.Structure):
"""A rectangle, used as part of detection results."""
_fields_ = [
('left', ctypes.c_int),
('top', ctypes.c_int),
('bottom', ctypes.c_int),
('right', ctypes.c_int),
]
class RectFC(ctypes.Structure):
"""A rectangle (operating on floats), used as part of processing options."""
_fields_ = [
('left', ctypes.c_float),
('top', ctypes.c_float),
('bottom', ctypes.c_float),
('right', ctypes.c_float),
]
class NormalizedRectC(ctypes.Structure):
"""A rectangle with rotation in normalized coordinates."""
_fields_ = [
('left', ctypes.c_float),
('top', ctypes.c_float),
('bottom', ctypes.c_float),
('right', ctypes.c_float),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/processors/__init__.py
+23
View File
@@ -0,0 +1,23 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks Components Processors API."""
import mediapipe.tasks.python.components.processors.classifier_options
ClassifierOptions = classifier_options.ClassifierOptions
# Remove unnecessary modules to avoid duplication in API docs.
del classifier_options
del mediapipe
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/processors/classifier_options.py
+64
View File
@@ -0,0 +1,64 @@
# Copyright 2022 The TensorFlow Authors.
#
# 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.
"""Classifier options data class."""
import dataclasses
from typing import Any, List, Optional
@dataclasses.dataclass
class ClassifierOptions:
"""Options for classification processor.
Attributes:
display_names_locale: The locale to use for display names specified through
the TFLite Model Metadata.
max_results: The maximum number of top-scored classification results to
return.
score_threshold: Overrides the ones provided in the model metadata. Results
below this value are rejected.
category_allowlist: Allowlist of category names. If non-empty, detection
results whose category name is not in this set will be filtered out.
Duplicate or unknown category names are ignored. Mutually exclusive with
`category_denylist`.
category_denylist: Denylist of category names. If non-empty, detection
results whose category name is in this set will be filtered out. Duplicate
or unknown category names are ignored. Mutually exclusive with
`category_allowlist`.
"""
display_names_locale: Optional[str] = None
max_results: Optional[int] = None
score_threshold: Optional[float] = None
category_allowlist: Optional[List[str]] = None
category_denylist: Optional[List[str]] = None
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, ClassifierOptions):
return False
return (
self.display_names_locale == other.display_names_locale
and self.max_results == other.max_results
and self.score_threshold == other.score_threshold
and self.category_allowlist == other.category_allowlist
and self.category_denylist == other.category_denylist
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/processors/classifier_options_c.py
+81
View File
@@ -0,0 +1,81 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""C types for ClassifierOptions."""
import ctypes
import mediapipe.tasks.python.components.processors.classifier_options as classifier_options_module
class ClassifierOptionsC(ctypes.Structure):
_fields_ = [
("display_names_locale", ctypes.c_char_p),
("max_results", ctypes.c_int),
("score_threshold", ctypes.c_float),
("category_allowlist", ctypes.POINTER(ctypes.c_char_p)),
("category_allowlist_count", ctypes.c_uint32),
("category_denylist", ctypes.POINTER(ctypes.c_char_p)),
("category_denylist_count", ctypes.c_uint32),
]
def convert_to_classifier_options_c(
src: classifier_options_module.ClassifierOptions,
) -> ClassifierOptionsC:
"""Converts a Python ClassifierOptions object to a ClassifierOptionsC object."""
options = ClassifierOptionsC()
if src.display_names_locale is not None:
options.display_names_locale = src.display_names_locale.encode("utf-8")
else:
options.display_names_locale = None
if src.max_results is not None:
options.max_results = src.max_results
else:
options.max_results = -1
if src.score_threshold is not None:
options.score_threshold = src.score_threshold
else:
options.score_threshold = 0.0
# Handle category_allowlist
if src.category_allowlist:
bytes_allowlist = [s.encode("utf-8") for s in src.category_allowlist]
allowlist_array_type = ctypes.c_char_p * len(bytes_allowlist)
allowlist_array = allowlist_array_type(*bytes_allowlist)
options.category_allowlist = ctypes.cast(
allowlist_array, ctypes.POINTER(ctypes.c_char_p)
)
options.category_allowlist_count = len(bytes_allowlist)
else:
options.category_allowlist = None
options.category_allowlist_count = 0
# Handle category_denylist
if src.category_denylist:
bytes_denylist = [s.encode("utf-8") for s in src.category_denylist]
denylist_array_type = ctypes.c_char_p * len(bytes_denylist)
denylist_array = denylist_array_type(*bytes_denylist)
options.category_denylist = ctypes.cast(
denylist_array, ctypes.POINTER(ctypes.c_char_p)
)
options.category_denylist_count = len(bytes_denylist)
else:
options.category_denylist = None
options.category_denylist_count = 0
return options
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/utils/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/components/utils/cosine_similarity.py
+68
View File
@@ -0,0 +1,68 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Cosine similarity utilities."""
import numpy as np
from mediapipe.tasks.python.components.containers import embedding_result
_Embedding = embedding_result.Embedding
def _compute_cosine_similarity(u: np.ndarray, v: np.ndarray):
"""Computes cosine similarity between two embeddings."""
if len(u) <= 0:
raise ValueError("Cannot compute cosing similarity on empty embeddings.")
norm_u = np.linalg.norm(u)
norm_v = np.linalg.norm(v)
if norm_u <= 0 or norm_v <= 0:
raise ValueError(
"Cannot compute cosine similarity on embedding with 0 norm.")
return u.dot(v) / (norm_u * norm_v)
def cosine_similarity(u: _Embedding, v: _Embedding) -> float:
"""Utility function to compute cosine similarity between two embedding.
May return an InvalidArgumentError if e.g. the feature vectors are of
different types (quantized vs. float), have different sizes, or have an
L2-norm of 0.
Args:
u: An embedding.
v: An embedding.
Returns:
Cosine similarity value.
"""
if len(u.embedding) != len(v.embedding):
raise ValueError(f"Cannot compute cosine similarity between embeddings "
f"of different sizes "
f"({len(u.embedding)} vs. {len(v.embedding)}).")
if u.embedding.dtype == float and v.embedding.dtype == float:
return _compute_cosine_similarity(u.embedding, v.embedding)
if u.embedding.dtype == np.uint8 and v.embedding.dtype == np.uint8:
return _compute_cosine_similarity(
u.embedding.view("int8").astype("float"),
v.embedding.view("int8").astype("float"),
)
raise ValueError("Cannot compute cosine similarity between quantized and "
"float embeddings.")
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/async_result_dispatcher.py
+232
View File
@@ -0,0 +1,232 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""A dispatcher to handle asynchronous results from a C library."""
import dataclasses
import enum
import logging
import os
import queue
import threading
from typing import Any, Callable, Optional, TypeVar, Union
from mediapipe.tasks.python.core import mediapipe_c_utils
CCallbackType = TypeVar("CCallbackType", bound=Callable[..., Any])
@dataclasses.dataclass(frozen=True)
class LiveStreamPacket:
"""A packet holding the contents of an async stream packet.
Attributes:
contents: The data that will be passed to the user-provided callback.
"""
contents: tuple[object, ...]
@dataclasses.dataclass(frozen=True)
class _ExceptionPacket:
"""A packet holding an exception during an async operation.
Attributes:
exception: The exception that will be logged in the dispatcher thread.
"""
exception: Exception
_AsyncResultPacket = Union[LiveStreamPacket, _ExceptionPacket]
class _DispatcherState(enum.Enum):
"""The state of the AsyncResultDispatcher."""
NOT_STARTED = 0
RUNNING = 1
SHUTTING_DOWN = 2
_PIPE_NOT_INITIALIZED = -1
class AsyncResultDispatcher:
"""A dispatcher to safely handle asynchronous results from a C library.
This class ensures that the user-provided Python callback is executed on a
dedicated, GIL-holding Python thread, rather than directly on the C library's
background thread. To use this class, you must supply a converter function
that converts the C data types and then pass the Python callback you want to
invoke to `wrap_callback()`, which returns a C callback that can be passed
to the C library.
The implementation uses OS level pipes to reduce the amount of Python code
that is run without the GIL lock. The dispatcher trhead is only started when
the C callback is first called.
Here's an example of how to use `AsyncResultDispatcher`:
```
def python_callback(result_data: MyResultType, timestamp: int):
print(f"Received result at {timestamp}: {result_data}")
def converter(c_result_ptr: ctypes.POINTER(CResult), c_timestamp: int):
py_result = MyResultType.from_ctypes(c_result_ptr[0])
py_timestamp = c_timestamp
return py_result, py_timestamp
dispatcher = AsyncResultDispatcher(converter=converter)
c_callback_type = ctypes.CFUNCTYPE(
None, # Return type (void)
ctypes.c_int32, # Status code
ctypes.POINTER(CResult), # First argument
ctypes.c_int64 # Second argument
)
c_callback = dispatcher.wrap_callback(python_callback, c_callback_type)
```
"""
_callback: Optional[Callable[..., None]] = None
_data_queue: queue.Queue[_AsyncResultPacket]
_pipe_read_fd: int = _PIPE_NOT_INITIALIZED
_pipe_write_fd: int = _PIPE_NOT_INITIALIZED
_state: _DispatcherState = _DispatcherState.NOT_STARTED
_dispatcher_thread: Optional[threading.Thread] = None
def __init__(self, converter: Callable[..., Any]):
"""Initializes the AsyncResultDispatcher.
Args:
converter: A function that takes the raw C-style arguments from the C
library callback and converts them into Python-friendly types.
"""
self._converter = converter
self._data_queue = queue.Queue()
def wrap_callback(
self,
python_callback: Optional[Callable[..., None]],
c_callback_type: type[CCallbackType],
) -> CCallbackType:
"""Returns a ctypes callback function that can be passed to a C library.
This function calls the converter function and schedules the user callback
for execution with the converted data. This function can only be called
once per instance.
Args:
python_callback: The user-defined result callback for processing live
stream data.
c_callback_type: The ctypes callback function wrapperof the function you
wish to create.
Raises:
RuntimeError: If a callback has already been assigned.
"""
if python_callback is None:
return c_callback_type()
if self._callback is not None:
raise RuntimeError("Callback already set.")
self._callback = python_callback
def c_callback(status_code: int, *args: Any) -> None:
"""The callback function that is invoked by the C library.
Args:
status_code: The status code that is passed to the C library callback.
*args: Any additional arguments passed to the C library callback.
"""
if status_code != 0:
exception = mediapipe_c_utils.convert_to_exception(status_code)
self._put_packet(_ExceptionPacket(exception=exception))
return
try:
py_args = self._converter(*args)
self._put_packet(LiveStreamPacket(contents=py_args))
except Exception as e: # pylint: disable=broad-except
self._put_packet(_ExceptionPacket(exception=e))
return c_callback_type(c_callback)
def _process_packet(self, packet: _AsyncResultPacket) -> None:
"""Processes a single packet from the data queue."""
if isinstance(packet, _ExceptionPacket):
logging.error("Error in async operation: %r", packet.exception)
elif isinstance(packet, LiveStreamPacket) and self._callback:
try:
self._callback(*packet.contents)
except Exception as e: # pylint: disable=broad-except
logging.exception("Error in callback: %r", e)
def _dispatcher_loop(self) -> None:
"""The main loop for the dispatcher thread."""
while True:
# Drain elements that are currently in the queue.
while not self._data_queue.empty():
packet = self._data_queue.get_nowait()
self._process_packet(packet)
if self._state == _DispatcherState.SHUTTING_DOWN:
break
try:
# Block until the C callback writes to the pipe again
os.read(self._pipe_read_fd, 1)
except OSError:
# Pipe might error during shutdown. Ignore.
continue
def _put_packet(self, packet: _AsyncResultPacket) -> None:
"""Puts a data packet into the queue and signals the dispatcher thread."""
if self._state == _DispatcherState.NOT_STARTED:
self._start()
elif self._state == _DispatcherState.SHUTTING_DOWN:
return
self._data_queue.put(packet)
try:
os.write(self._pipe_write_fd, b"\0")
except OSError:
# Pipe might error during shutdown. Ignore.
pass
def close(self) -> None:
"""Shuts down the dispatcher thread and cleans up resources."""
if self._state == _DispatcherState.RUNNING:
self._state = _DispatcherState.SHUTTING_DOWN
try:
# Write a final byte to unblock the os.read() call.
os.write(self._pipe_write_fd, b"\0")
except OSError:
pass
if self._dispatcher_thread:
self._dispatcher_thread.join()
os.close(self._pipe_read_fd)
os.close(self._pipe_write_fd)
def _start(self) -> None:
"""Starts the dispatcher thread."""
if self._callback is None:
raise ValueError("Dispatcher should not be started without a callback.")
self._state = _DispatcherState.RUNNING
self._pipe_read_fd, self._pipe_write_fd = os.pipe()
self._dispatcher_thread = threading.Thread(
target=self._dispatcher_loop, daemon=True
)
self._dispatcher_thread.start()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/base_options.py
+125
View File
@@ -0,0 +1,125 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Base options for MediaPipe Task APIs."""
import dataclasses
import enum
import platform
import sys
import types
from typing import Any, Optional
import certifi
from mediapipe.tasks.python.core import base_options_c as base_options_c_lib
from mediapipe.tasks.python.core.optional_dependencies import doc_controls
# C enum value for host environment.
HOST_ENVIRONMENT_PYTHON: int = 3
class _HostSystem(enum.IntEnum):
"""An mapping for the C enum values for the host operating system.
Attributes:
HOST_SYSTEM_UNKNOWN: Unknown host system.
HOST_SYSTEM_LINUX: Linux host system.
HOST_SYSTEM_MAC: MacOS host system.
HOST_SYSTEM_WINDOWS: Windows host system.
"""
HOST_SYSTEM_UNKNOWN = 0
HOST_SYSTEM_LINUX = 1
HOST_SYSTEM_MAC = 2
HOST_SYSTEM_WINDOWS = 3
_HOST_SYSTEM_BY_PLATFORM = types.MappingProxyType({
'Darwin': _HostSystem.HOST_SYSTEM_MAC,
'Windows': _HostSystem.HOST_SYSTEM_WINDOWS,
'Linux': _HostSystem.HOST_SYSTEM_LINUX,
})
@dataclasses.dataclass
class BaseOptions:
"""Base options for MediaPipe Tasks' Python APIs.
Represents external model asset used by the Task APIs. The files can be
specified by one of the following two ways:
(1) model asset file path in `model_asset_path`.
(2) model asset contents loaded in `model_asset_buffer`.
If more than one field of these fields is provided, they are used in this
precedence order.
Attributes:
model_asset_path: Path to the model asset file.
model_asset_buffer: The model asset file contents as bytes.
delegate: Acceleration to use. Supported values are GPU and CPU. GPU support
is currently limited to Ubuntu platforms.
"""
class Delegate(enum.IntEnum):
CPU = 0
GPU = 1
model_asset_path: Optional[str] = None
model_asset_buffer: Optional[bytes] = None
delegate: Optional[Delegate] = None
@doc_controls.do_not_generate_docs
def to_ctypes(self) -> base_options_c_lib.BaseOptionsC:
"""Creates a BaseOptionsC struct from the BaseOptions object."""
ca_bundle_path = certifi.where()
host_system = _HOST_SYSTEM_BY_PLATFORM.get(
platform.system(), _HostSystem.HOST_SYSTEM_UNKNOWN
)
options = base_options_c_lib.BaseOptionsC()
options.model_asset_buffer = self.model_asset_buffer
options.model_asset_buffer_count = (
len(self.model_asset_buffer) if self.model_asset_buffer else 0
)
options.model_asset_path = (
self.model_asset_path.encode('utf-8') if self.model_asset_path else None
)
options.delegate = self.delegate.value if self.delegate else 0
options.host_environment = HOST_ENVIRONMENT_PYTHON
host_version, *_ = sys.version.split(' ')
options.host_system = host_system
options.host_version = host_version.encode('utf-8')
options.ca_bundle_path = ca_bundle_path.encode('utf-8')
return options
def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object.
Args:
other: The object to be compared with.
Returns:
True if the objects are equal.
"""
if not isinstance(other, BaseOptions):
return False
return (
self.model_asset_path == other.model_asset_path
and self.model_asset_buffer == other.model_asset_buffer
and self.delegate == other.delegate
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/base_options_c.py
+42
View File
@@ -0,0 +1,42 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""C types for BaseOptions."""
import ctypes
class BaseOptionsC(ctypes.Structure):
"""C types for BaseOptions.
Attributes:
model_asset_buffer: `bytes`, the model asset buffer.
model_asset_buffer_count: `int`, the size of the model asset buffer.
model_asset_path: `bytes`, the path to the model asset.
delegate: `int`, the delegate to use.
host_environment: `int`, the environment in which the task is running.
host_system: `int`, the system on which the task is running.
host_version: `bytes`, the Python version as a UTF-8 string.
ca_bundle_path: `bytes`, the path to the CA bundle file as a UTF-8 string.
"""
_fields_ = [
("model_asset_buffer", ctypes.c_char_p),
("model_asset_buffer_count", ctypes.c_uint),
("model_asset_path", ctypes.c_char_p),
("delegate", ctypes.c_int),
("host_environment", ctypes.c_int),
("host_system", ctypes.c_int),
("host_version", ctypes.c_char_p),
("ca_bundle_path", ctypes.c_char_p),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/mediapipe_c_bindings.py
+105
View File
@@ -0,0 +1,105 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""MediaPipe Text tasks shared library."""
import ctypes
import os
import platform
from typing import Any, List, Optional, Sequence
from importlib import resources
from mediapipe.tasks.python.core import mediapipe_c_utils
from mediapipe.tasks.python.core import serial_dispatcher
_BASE_LIB_PATH = 'mediapipe/tasks/c/'
_shared_lib = None
_CFunction = mediapipe_c_utils.CFunction
def load_raw_library(signatures: Sequence[_CFunction] = ()) -> ctypes.CDLL:
"""Loads the raw ctypes.CDLL shared library and registers signatures.
This function loads the raw ctypes.CDLL shared library if it hasn't been
loaded yet, otherwise it re-uses the existing instance. It attaches the
provided signatures and makes them available across all callsites.
Args:
signatures: The ctypes function signatures to register in the library.
Returns:
The ctypes shared library.
"""
global _shared_lib
if _shared_lib is None:
if os.name == 'posix':
if platform.system() == 'Darwin': # macOS
lib_filename = 'libmediapipe.dylib'
else: # Linux
lib_filename = 'libmediapipe.so'
else: # Windows
lib_filename = 'libmediapipe.dll'
lib_path_context = resources.files('mediapipe.tasks.c')
absolute_lib_path = str(lib_path_context / lib_filename)
_shared_lib = ctypes.CDLL(absolute_lib_path)
for signature in signatures:
c_func = getattr(_shared_lib, signature.func_name)
c_func.argtypes = signature.argtypes
c_func.restype = signature.restype
# Register "MpErrorFree()"
_shared_lib.MpErrorFree.argtypes = [ctypes.c_void_p]
_shared_lib.MpErrorFree.restype = None
return _shared_lib
def load_shared_library(
signatures: Sequence[_CFunction] = (),
) -> serial_dispatcher.SerialDispatcher:
"""Loads the shared library in a SerialDispatcher and registers signatures.
This function creates a thread-safe wrapper for the shared library that
dispatches all calls through a single dedicated thread. Every call to this
function uses a different dispatch library with its own thread.
Args:
signatures: The ctypes function signatures to register in the library.
Returns:
A thread-safe wrapper for the ctypes shared library.
"""
raw_lib = load_raw_library()
return serial_dispatcher.SerialDispatcher(raw_lib, signatures)
def convert_strings_to_ctypes_array(
str_list: Optional[List[str]],
) -> Optional[Any]:
"""Converts a list of Python strings to a ctypes array of c_char_p.
Args:
str_list: A list of strings, or None.
Returns:
A ctypes array of c_char_p, or None if the input is None or empty.
"""
if not str_list: # Handles None and empty list
return None
num_elements = len(str_list)
c_array = (ctypes.c_char_p * num_elements)()
encoded_strings = [s.encode('utf-8') for s in str_list]
c_array[:] = encoded_strings
return c_array
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/mediapipe_c_utils.py
+232
View File
@@ -0,0 +1,232 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""Common MediaPipe CTypes definitions and utilities."""
import atexit
import concurrent.futures
import ctypes
import dataclasses
import enum
import functools
import threading
from typing import Any, Callable, Optional, Sequence, Union
# Global shutdown state to ignore calls during Python shutdown. This is used to
# prevent forwarding calls to the underlying executor, which prints warning logs
# to the console after shutdown.
_shutdown = False
_global_shutdown_lock = threading.Lock()
def _python_exit():
global _shutdown
with _global_shutdown_lock:
_shutdown = True
def is_shutdown():
with _global_shutdown_lock:
return _shutdown
atexit.register(_python_exit)
@dataclasses.dataclass(frozen=True)
class CFunction:
"""Stores the ctypes signature information for a C function.
Attributes:
func_name: The name of the C function.
argtypes: The ctypes types for the function arguments.
restype: The ctypes type for the function's return value.
"""
func_name: str
argtypes: list[Any]
restype: Any
def create_python_wrapper(
self, lib: ctypes.CDLL, executor: concurrent.futures.Executor
) -> Callable[..., Any]:
"""Creates a callable wrapper that executes the C function on the executor.
Args:
lib: The ctypes shared library.
executor: The executor to use for the C function.
Returns:
A callable that executes the C function on the executor.
"""
c_func = getattr(lib, self.func_name)
c_func.argtypes = self.argtypes
c_func.restype = self.restype
@functools.wraps(c_func)
def dispatcher_wrapper(*args, **kwargs) -> Any:
future = executor.submit(c_func, *args, **kwargs)
return future.result()
return dispatcher_wrapper
@dataclasses.dataclass(frozen=True)
class CStatusFunction(CFunction):
"""Stores the ctypes signature for a status-returning C function.
The functions is required to take an error message argument as the last
argument and return an `MpStatus' integer code.
"""
def __init__(self, func_name: str, core_argtypes: Sequence[Any]):
"""Initializes the CStatusFunction.
Args:
func_name: The name of the C function.
core_argtypes: The ctypes types for the function arguments excluding the
error message argument.
"""
full_argtypes = list(core_argtypes) + [ctypes.POINTER(ctypes.c_char_p)]
super().__init__(func_name, full_argtypes, restype=ctypes.c_int)
def create_python_wrapper(
self, lib: ctypes.CDLL, executor: concurrent.futures.Executor
) -> Callable[..., None]:
"""Creates a callable wrapper that executes the C function on the executor.
This function raises a Python exception based on the MpStatus code returned
by the C API function and the error message argument.
Args:
lib: The ctypes shared library.
executor: The executor to use for the C function.
Returns:
A callable that executes the C function on the executor.
"""
c_func = getattr(lib, self.func_name)
c_func.argtypes = self.argtypes
c_func.restype = self.restype
error_msg_index = len(self.argtypes)
@functools.wraps(c_func)
def dispatcher_wrapper(*args, **kwargs) -> None:
if is_shutdown():
return
def dispatch_and_free():
# Add the error message argument to the list of arguments.
error_msg = ctypes.c_char_p()
call_args = list(args)
call_args.insert(error_msg_index, ctypes.byref(error_msg))
status = c_func(*call_args, **kwargs)
try:
handle_status(status, error_msg)
finally:
if error_msg.value is not None:
lib.MpErrorFree(error_msg)
executor.submit(dispatch_and_free).result()
return dispatcher_wrapper
class MpStatus(enum.IntEnum):
"""Status codes for MediaPipe C API functions."""
MP_OK = 0
MP_CANCELLED = 1
MP_UNKNOWN = 2
MP_INVALID_ARGUMENT = 3
MP_DEADLINE_EXCEEDED = 4
MP_NOT_FOUND = 5
MP_ALREADY_EXISTS = 6
MP_PERMISSION_DENIED = 7
MP_RESOURCE_EXHAUSTED = 8
MP_FAILED_PRECONDITION = 9
MP_ABORTED = 10
MP_OUT_OF_RANGE = 11
MP_UNIMPLEMENTED = 12
MP_INTERNAL = 13
MP_UNAVAILABLE = 14
MP_DATA_LOSS = 15
MP_UNAUTHENTICATED = 16
def convert_to_exception(
status: int, error_message: Optional[ctypes.c_char_p] = None
) -> Union[Exception, None]:
"""Returns an exception based on the MpStatus code, or None if MP_OK."""
error_str = (
error_message.value.decode('utf-8')
if error_message and error_message.value is not None
else None
)
if status == MpStatus.MP_OK:
return None
elif status == MpStatus.MP_CANCELLED:
return TimeoutError(error_str or 'Cancelled')
elif status == MpStatus.MP_UNKNOWN:
return RuntimeError(error_str or 'Unknown error')
elif status == MpStatus.MP_INVALID_ARGUMENT:
return ValueError(error_str or 'Invalid argument')
elif status == MpStatus.MP_DEADLINE_EXCEEDED:
return TimeoutError(error_str or 'Deadline exceeded')
elif status == MpStatus.MP_NOT_FOUND:
return FileNotFoundError(error_str or 'Not found')
elif status == MpStatus.MP_ALREADY_EXISTS:
return FileExistsError(error_str or 'Already exists')
elif status == MpStatus.MP_PERMISSION_DENIED:
return PermissionError(error_str or 'Permission denied')
elif status == MpStatus.MP_RESOURCE_EXHAUSTED:
return RuntimeError(error_str or 'Resource exhausted')
elif status == MpStatus.MP_FAILED_PRECONDITION:
return RuntimeError(error_str or 'Failed precondition')
elif status == MpStatus.MP_ABORTED:
return RuntimeError(error_str or 'Aborted')
elif status == MpStatus.MP_OUT_OF_RANGE:
return IndexError(error_str or 'Out of range')
elif status == MpStatus.MP_UNIMPLEMENTED:
return NotImplementedError(error_str or 'Unimplemented')
elif status == MpStatus.MP_INTERNAL:
return RuntimeError(error_str or 'Internal error')
elif status == MpStatus.MP_UNAVAILABLE:
return RuntimeError(error_str or 'Unavailable')
elif status == MpStatus.MP_DATA_LOSS:
return RuntimeError(error_str or 'Data loss')
elif status == MpStatus.MP_UNAUTHENTICATED:
return PermissionError(error_str or 'Unauthenticated')
else:
return RuntimeError(error_str or f'Unexpected status: {status}')
def handle_status(
status: int, error_message: Optional[ctypes.c_char_p] = None
) -> None:
"""Checks the MpStatus and raises an error if not kMpOk.
Args:
status: The MpStatus code to check.
error_message: The error message to use if an exception is raised. Uses a
default error message if None.
Raises:
An error if the MpStatus is not kMpOk.
"""
exception = convert_to_exception(status, error_message)
if exception:
raise exception
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/optional_dependencies.py
+25
View File
@@ -0,0 +1,25 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks' common but optional dependencies."""
# TensorFlow isn't a dependency of mediapipe pip package. It's only
# required in the API docgen pipeline so we'll ignore it if tensorflow is not
# installed.
try:
from tensorflow.tools.docs import doc_controls
except ModuleNotFoundError:
# Replace the real doc_controls.do_not_generate_docs with an no-op
doc_controls = lambda: None
no_op = lambda x: x
setattr(doc_controls, 'do_not_generate_docs', no_op)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/core/serial_dispatcher.py
+91
View File
@@ -0,0 +1,91 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""A thread-safe dispatcher for a ctypes."""
import concurrent.futures
import ctypes
import threading
from typing import Any, Sequence
from mediapipe.tasks.python.core import mediapipe_c_utils
class SerialDispatcher:
"""A wrapper class for a ctypes.CDLL object that serializes all calls.
This ensures that functions from a non-thread-safe C library are called
sequentially from a single dedicated thread, preventing race conditions
and segmentation faults.
If a function is a CStatusFunction, the dispatcher will raise a Python
exception if the returned MpStatus code is not kMpOk.
"""
# Enable dynamic attributes as we register methods on this class via
# reflection. This allows these methods to be used with type checking.
_HAS_DYNAMIC_ATTRIBUTES = True
def __init__(
self, lib: ctypes.CDLL, signatures: Sequence[mediapipe_c_utils.CFunction]
):
"""Initializes the SerialDispatcher.
Args:
lib: The ctypes.CDLL object to wrap.
signatures: the CFunction objects specifying the functions to wrap.
"""
self._lib = lib
self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
self._lock = threading.Lock()
self._is_closed = False
for signature in signatures:
self._register_signature(signature)
def _register_signature(self, signature: mediapipe_c_utils.CFunction):
"""Registers a wrapped C function as a method on the SerialDispatcher.
This method attaches a wrapper method to this class that dispatches all
calls to the underlying C function through this class' executor. The
methods are directly attached to the SerialDispatcher object, allowing them
to be used transparently.
Args:
signature: The CFunction object specifying the C function to wrap and
register.
"""
handler = signature.create_python_wrapper(self._lib, self._executor)
def shutdown_aware_handler(*args, **kwargs) -> Any:
with self._lock:
if self._is_closed:
return
return handler(*args, **kwargs)
setattr(self, signature.func_name, shutdown_aware_handler)
def close(self):
"""Shuts down the dispatcher and waits for pending tasks to complete."""
with self._lock:
if self._is_closed:
return
self._is_closed = True
self._executor.shutdown(wait=True)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
del exc_type, exc_val, exc_tb
self.close()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/__init__.py
+14
View File
@@ -0,0 +1,14 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""MediaPipe GenAI Python API."""
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/__init__.py
+23
View File
@@ -0,0 +1,23 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks GenAI Bundler API."""
import mediapipe.tasks.python.genai.bundler.llm_bundler
BundleConfig = llm_bundler.BundleConfig
create_bundle = llm_bundler.create_bundle
# Remove unnecessary modules to avoid duplication in API docs.
del llm_bundler
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler.py
+236
View File
@@ -0,0 +1,236 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Functions to perform llm packing."""
import ctypes
import dataclasses
import enum
from typing import List, Optional
from mediapipe.tasks.python.core import mediapipe_c_bindings
from mediapipe.tasks.python.core import mediapipe_c_utils
from mediapipe.tasks.python.core import serial_dispatcher
from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options
from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options_c
from mediapipe.tasks.python.genai.converter import external_dependencies
from mediapipe.tasks.python.metadata.metadata_writers import model_asset_bundle_utils
sentencepiece = external_dependencies.sentencepiece
_CTYPES_SIGNATURES = (
mediapipe_c_utils.CFunction(
func_name="MpLlmBundlerGenerateMetadata",
argtypes=[
ctypes.POINTER(
llm_bundler_metadata_options_c.LlmBundlerMetadataOptionsC
),
ctypes.POINTER(ctypes.c_int),
],
restype=ctypes.POINTER(ctypes.c_char),
),
mediapipe_c_utils.CFunction(
func_name="MpLlmBundlerFreeMetadata",
argtypes=[
ctypes.POINTER(ctypes.c_char),
],
restype=None,
),
)
@dataclasses.dataclass(frozen=True)
class BundleConfig:
"""Config for LLM Bundler.
Attributes:
tflite_model: Path to the multi-signature tflite model with "prefill" and
"decode" signatures converted using ODML Transformers APIs.
tokenizer_model: Path to the tokenizer model. Currently only SentencePience
tokenizer is supported. As such, tokenizer.model proto is expected to be
passed here.
start_token: Token that will be used to signify the beginning of a sequence.
stop_tokens: Tokens that will be used to signify the end of a sequence.
output_filename: Name of the generated `.task` file containing the Bundle.
enable_bytes_to_unicode_mapping: Enables GPT-2 style bytes to unicode
mapping. For more details see:
https://github.com/openai/gpt-2/blob/master/src/encoder.py#L9
system_prompt: The starting text would be feed to the model on the start of
each session, commonly called the System prompt. This is useful for
instruction tuned models and pre-conditioning the model behavior. This
field is mutually exclusive with prompt_prefix_system and
prompt_suffix_system.
prompt_prefix_user: The prefix that should be prepended to each user portion
of prompts passed to the model.
prompt_suffix_user: The suffix that should be appended to each user portion
of prompts passed to the model.
prompt_prefix_model: The prefix that should be prepended to each model
portion of prompts passed to the model.
prompt_suffix_model: The suffix that should be appended to each model
portion of prompts passsed to the model.
prompt_prefix_system: The prefix that should be prepended to each system
portion of prompts passed to the model.
prompt_suffix_system: The suffix that should be appended to each system
portion of prompts passsed to the model.
user_role_token: The token that indicates the user's turn.
system_role_token: The token that indicates the system's turn.
model_role_token: The token that indicates the model's turn.
end_role_token: The token that indicates the end of the prompt.
tflite_embedder: Path to the embedding model if the embedding lookup is done
externally.
tflite_per_layer_embedder: Path to the per layer embedding model if the per
layer embedding lookup is done externally.
tflite_vision_encoder: Path to the vision encoder model.
tflite_vision_adapter: Path to the vision adapter model.
"""
tflite_model: str
tokenizer_model: str
start_token: str
stop_tokens: List[str]
output_filename: str
enable_bytes_to_unicode_mapping: bool = False
system_prompt: Optional[str] = None
prompt_prefix_user: Optional[str] = None
prompt_suffix_user: Optional[str] = None
prompt_prefix_model: Optional[str] = None
prompt_suffix_model: Optional[str] = None
prompt_prefix_system: Optional[str] = None
prompt_suffix_system: Optional[str] = None
user_role_token: Optional[str] = None
system_role_token: Optional[str] = None
model_role_token: Optional[str] = None
end_role_token: Optional[str] = None
tflite_embedder: Optional[str] = None
tflite_per_layer_embedder: Optional[str] = None
tflite_vision_encoder: Optional[str] = None
tflite_vision_adapter: Optional[str] = None
class _BundleTags(enum.Enum):
"""Bundle tags."""
TF_LITE_PREFILL_DECODE = 1
TOKENIZER_MODEL = 2
METADATA = 3
TF_LITE_EMBEDDER = 4
TF_LITE_PER_LAYER_EMBEDDER = 5
TF_LITE_VISION_ENCODER = 6
TF_LITE_VISION_ADAPTER = 7
def _validate_config(config: BundleConfig):
"""Validates the given config.
Args:
config: The config to validate.
Raises:
ValueError if the config is invalid. Invalid configs can be:
- tokenizer_model is not a valid SentencePiece model.
- start_token is not a valid token in the tokenizer_model.
- stop_tokens is not a list of valid tokens in the tokenizer_model.
- system_prompt and prompt_*_system are both set.
"""
if not isinstance(config.stop_tokens, list):
raise ValueError("stop_tokens must be a list of strings.")
if not config.stop_tokens:
raise ValueError("stop_tokens must be non-empty.")
try:
sp = sentencepiece.SentencePieceProcessor()
sp.Load(config.tokenizer_model)
except Exception as e:
raise ValueError(
f"Failed to load tokenizer model from {config.tokenizer_model}. "
"Please ensure you are passing a valid SentencePiece model."
) from e
if config.system_prompt and (
config.prompt_prefix_system or config.prompt_suffix_system
):
raise ValueError(
"system_prompt and prompt_*_system are mutually exclusive."
)
def create_bundle(config: BundleConfig):
"""Creates a bundle from the given config."""
_validate_config(config)
artifacts = {}
with open(config.tflite_model, "rb") as f:
artifacts[_BundleTags.TF_LITE_PREFILL_DECODE.name] = f.read()
with open(config.tokenizer_model, "rb") as f:
artifacts[_BundleTags.TOKENIZER_MODEL.name] = f.read()
lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES)
with serial_dispatcher.SerialDispatcher(
lib, _CTYPES_SIGNATURES
) as bundler_lib:
metadata_options = llm_bundler_metadata_options.LlmBundlerMetadataOptions(
start_token=config.start_token,
stop_tokens=config.stop_tokens,
enable_bytes_to_unicode_mapping=config.enable_bytes_to_unicode_mapping,
system_prompt=config.system_prompt,
prompt_prefix_user=config.prompt_prefix_user,
prompt_suffix_user=config.prompt_suffix_user,
prompt_prefix_model=config.prompt_prefix_model,
prompt_suffix_model=config.prompt_suffix_model,
prompt_prefix_system=config.prompt_prefix_system,
prompt_suffix_system=config.prompt_suffix_system,
user_role_token=config.user_role_token,
system_role_token=config.system_role_token,
model_role_token=config.model_role_token,
end_role_token=config.end_role_token,
)
metadata_options_c = metadata_options.to_ctypes()
buffer_size = ctypes.c_int()
metadata_buffer = None
try:
metadata_buffer = bundler_lib.MpLlmBundlerGenerateMetadata(
ctypes.byref(metadata_options_c), ctypes.byref(buffer_size)
)
artifacts[_BundleTags.METADATA.name] = ctypes.string_at(
metadata_buffer, buffer_size.value
)
finally:
if metadata_buffer:
bundler_lib.MpLlmBundlerFreeMetadata(metadata_buffer)
if config.tflite_embedder:
with open(config.tflite_embedder, "rb") as f:
artifacts[_BundleTags.TF_LITE_EMBEDDER.name] = f.read()
if config.tflite_per_layer_embedder:
with open(config.tflite_per_layer_embedder, "rb") as f:
artifacts[_BundleTags.TF_LITE_PER_LAYER_EMBEDDER.name] = f.read()
if config.tflite_vision_encoder:
with open(config.tflite_vision_encoder, "rb") as f:
artifacts[_BundleTags.TF_LITE_VISION_ENCODER.name] = f.read()
if config.tflite_vision_adapter:
with open(config.tflite_vision_adapter, "rb") as f:
artifacts[_BundleTags.TF_LITE_VISION_ADAPTER.name] = f.read()
output_filename = config.output_filename
if not output_filename.endswith(".task"):
output_filename = config.output_filename + ".task"
model_asset_bundle_utils.create_model_asset_bundle(
artifacts,
output_filename,
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler_metadata_options.py
+142
View File
@@ -0,0 +1,142 @@
# Copyright 2026 The MediaPipe Authors.
#
# 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.
"""LlmBundlerMetadataOptions Python dataclass."""
import ctypes
import dataclasses
from typing import Optional
from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options_c
@dataclasses.dataclass
class LlmBundlerMetadataOptions:
"""Options for bundling LLM metadata.
Attributes:
start_token: Token that will be used to signify the beginning of a sequence.
stop_tokens: Tokens that will be used to signify the end of a sequence.
enable_bytes_to_unicode_mapping: Enables GPT-2 style bytes to unicode
mapping. For more details see:
https://github.com/openai/gpt-2/blob/master/src/encoder.py#L9
system_prompt: The starting text would be feed to the model on the start of
each session, commonly called the System prompt. This is useful for
instruction tuned models and pre-conditioning the model behavior. This
field is mutually exclusive with prompt_prefix_system and
prompt_suffix_system.
prompt_prefix_user: The prefix that should be prepended to each user portion
of prompts passed to the model.
prompt_suffix_user: The suffix that should be appended to each user portion
of prompts passed to the model.
prompt_prefix_model: The prefix that should be prepended to each model
portion of prompts passed to the model.
prompt_suffix_model: The suffix that should be appended to each model
portion of prompts passsed to the model.
prompt_prefix_system: The prefix that should be prepended to each system
portion of prompts passed to the model.
prompt_suffix_system: The suffix that should be appended to each system
portion of prompts passsed to the model.
user_role_token: The token that indicates the user's turn.
system_role_token: The token that indicates the system's turn.
model_role_token: The token that indicates the model's turn.
end_role_token: The token that indicates the end of the prompt.
"""
start_token: str
stop_tokens: list[str]
enable_bytes_to_unicode_mapping: bool = False
system_prompt: Optional[str] = None
prompt_prefix_user: Optional[str] = None
prompt_suffix_user: Optional[str] = None
prompt_prefix_model: Optional[str] = None
prompt_suffix_model: Optional[str] = None
prompt_prefix_system: Optional[str] = None
prompt_suffix_system: Optional[str] = None
user_role_token: Optional[str] = None
system_role_token: Optional[str] = None
model_role_token: Optional[str] = None
end_role_token: Optional[str] = None
def to_ctypes(
self,
) -> llm_bundler_metadata_options_c.LlmBundlerMetadataOptionsC:
"""Converts the object to a ctypes structure."""
stop_tokens_c = (ctypes.c_char_p * len(self.stop_tokens))()
for i, token in enumerate(self.stop_tokens):
stop_tokens_c[i] = token.encode("utf-8")
start_token_c = self.start_token.encode("utf-8")
system_prompt_c = (
self.system_prompt.encode("utf-8") if self.system_prompt else None
)
prompt_prefix_user_c = (
self.prompt_prefix_user.encode("utf-8")
if self.prompt_prefix_user
else None
)
prompt_suffix_user_c = (
self.prompt_suffix_user.encode("utf-8")
if self.prompt_suffix_user
else None
)
prompt_prefix_model_c = (
self.prompt_prefix_model.encode("utf-8")
if self.prompt_prefix_model
else None
)
prompt_suffix_model_c = (
self.prompt_suffix_model.encode("utf-8")
if self.prompt_suffix_model
else None
)
prompt_prefix_system_c = (
self.prompt_prefix_system.encode("utf-8")
if self.prompt_prefix_system
else None
)
prompt_suffix_system_c = (
self.prompt_suffix_system.encode("utf-8")
if self.prompt_suffix_system
else None
)
user_role_token_c = (
self.user_role_token.encode("utf-8") if self.user_role_token else None
)
system_role_token_c = (
self.system_role_token.encode("utf-8")
if self.system_role_token
else None
)
model_role_token_c = (
self.model_role_token.encode("utf-8") if self.model_role_token else None
)
end_role_token_c = (
self.end_role_token.encode("utf-8") if self.end_role_token else None
)
return llm_bundler_metadata_options_c.LlmBundlerMetadataOptionsC(
start_token=start_token_c,
stop_tokens=stop_tokens_c,
num_stop_tokens=len(self.stop_tokens),
enable_bytes_to_unicode_mapping=self.enable_bytes_to_unicode_mapping,
system_prompt=system_prompt_c,
prompt_prefix_user=prompt_prefix_user_c,
prompt_suffix_user=prompt_suffix_user_c,
prompt_prefix_model=prompt_prefix_model_c,
prompt_suffix_model=prompt_suffix_model_c,
prompt_prefix_system=prompt_prefix_system_c,
prompt_suffix_system=prompt_suffix_system_c,
user_role_token=user_role_token_c,
system_role_token=system_role_token_c,
model_role_token=model_role_token_c,
end_role_token=end_role_token_c,
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler_metadata_options_c.py
+38
View File
@@ -0,0 +1,38 @@
# Copyright 2026 The MediaPipe Authors.
#
# 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.
"""Ctypes wrapper for LlmBundlerMetadataOptions."""
import ctypes
class LlmBundlerMetadataOptionsC(ctypes.Structure):
"""Wrapper for LlmBundlerMetadataOptions in C."""
_fields_ = [
("start_token", ctypes.c_char_p),
("stop_tokens", ctypes.POINTER(ctypes.c_char_p)),
("num_stop_tokens", ctypes.c_int),
("enable_bytes_to_unicode_mapping", ctypes.c_bool),
("system_prompt", ctypes.c_char_p),
("prompt_prefix_user", ctypes.c_char_p),
("prompt_suffix_user", ctypes.c_char_p),
("prompt_prefix_model", ctypes.c_char_p),
("prompt_suffix_model", ctypes.c_char_p),
("prompt_prefix_system", ctypes.c_char_p),
("prompt_suffix_system", ctypes.c_char_p),
("user_role_token", ctypes.c_char_p),
("system_role_token", ctypes.c_char_p),
("model_role_token", ctypes.c_char_p),
("end_role_token", ctypes.c_char_p),
]
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/bundler/llm_bundler_metadata_options_test.py
+60
View File
@@ -0,0 +1,60 @@
# Copyright 2026 The MediaPipe Authors.
#
# 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.
from absl.testing import absltest
from mediapipe.tasks.python.genai.bundler import llm_bundler_metadata_options
class LlmBundlerMetadataOptionsTest(absltest.TestCase):
def test_to_ctypes(self):
options = llm_bundler_metadata_options.LlmBundlerMetadataOptions(
start_token="start",
stop_tokens=["stop1", "stop2"],
enable_bytes_to_unicode_mapping=True,
system_prompt="system",
prompt_prefix_user="user_prefix",
prompt_suffix_user="user_suffix",
prompt_prefix_model="model_prefix",
prompt_suffix_model="model_suffix",
prompt_prefix_system=None,
prompt_suffix_system=None,
user_role_token="user_role",
system_role_token=None,
model_role_token="model_role",
end_role_token="end_role",
)
options_c = options.to_ctypes()
self.assertEqual(options_c.start_token, b"start")
self.assertEqual(options_c.stop_tokens[0], b"stop1")
self.assertEqual(options_c.stop_tokens[1], b"stop2")
self.assertEqual(options_c.num_stop_tokens, 2)
self.assertTrue(options_c.enable_bytes_to_unicode_mapping)
self.assertEqual(options_c.system_prompt, b"system")
self.assertEqual(options_c.prompt_prefix_user, b"user_prefix")
self.assertEqual(options_c.prompt_suffix_user, b"user_suffix")
self.assertEqual(options_c.prompt_prefix_model, b"model_prefix")
self.assertEqual(options_c.prompt_suffix_model, b"model_suffix")
self.assertIsNone(options_c.prompt_prefix_system)
self.assertIsNone(options_c.prompt_suffix_system)
self.assertEqual(options_c.user_role_token, b"user_role")
self.assertIsNone(options_c.system_role_token)
self.assertEqual(options_c.model_role_token, b"model_role")
self.assertEqual(options_c.end_role_token, b"end_role")
if __name__ == "__main__":
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/__init__.py
+24
View File
@@ -0,0 +1,24 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""MediaPipe Tasks GenAI Converter API."""
import mediapipe.tasks.python.genai.converter.llm_converter
ConversionConfig = llm_converter.ConversionConfig
convert_checkpoint = llm_converter.convert_checkpoint
# Remove unnecessary modules to avoid duplication in API docs.
del mediapipe
del llm_converter
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/converter_base.py
+184
View File
@@ -0,0 +1,184 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Defines a couple base classes for the conversion/quantization process."""
from typing import Iterator
import os
from typing import Dict, List, Optional, Tuple
import numpy as np
class QuantizationAction:
"""Container of the tensor values and its corresponding quantization settings.
The container is responsible for hosting all of the information that is
required to execute the weight-only quantization.
Attributes:
tensor_name: A string that represents the input tensor name.
tensor_value: A numpy array that contains the unquantized tensor values.
target_name: A string that represents the updated tensor name.
quantize_axis: A list of integers representing the dimensions to be
quantized along. For example, if an input tensor has shape [128, 256] and
the quantize_axis==[0], it means the quantization happens along the 0-th
dimension, resulting in [256] scaling factors.
quantize_bits: An integer that specifies the target quantization bits. It
currently only supports either 8 or 4 bits.
pack_dim: An integer specifying which dimension to pack the quantized bits.
This is only applicable when the quantize_bits == 4.
is_quantized: A boolean indicating whether the tensor is already quantized
in the checkpoint. Default to False.
"""
def __init__(
self,
tensor_name: str,
tensor_value: Optional[np.ndarray] = None,
target_name: Optional[str] = None,
quantize_axis: Optional[List[int]] = None,
quantize_bits: Optional[int] = None,
pack_dim: Optional[int] = 0,
is_quantized: bool = False,
):
"""Initializes the model attributes."""
self.tensor_name = tensor_name
self.tensor_value = tensor_value
self.target_name = target_name
self.quantize_axis = quantize_axis
self.quantize_bits = quantize_bits
self.pack_dim = pack_dim
self.is_quantized = is_quantized
def __str__(self) -> str:
output_string = "QuantizationAction(\n"
output_string += f" tensor_name: {self.tensor_name}\n"
output_string += f" target_name: {self.target_name}\n"
output_string += f" quantize_axis: {self.quantize_axis}\n"
output_string += f" quantize_bits: {self.quantize_bits}\n"
output_string += f" pack_dim: {self.pack_dim}\n"
output_string += f" is_quantized: {self.is_quantized}\n"
if self.tensor_value is not None:
output_string += f" tensor_value: {self.tensor_value.shape}\n"
output_string += ")\n"
return output_string
class CkptLoaderBase:
"""Base class for loading the checkpoint.
This class is responsible for loading the checkpoint files into the layer
weight tensors (as numpy arrays) + quantization setting information (8/4
bits). The returned data should be a list of QuantizationAction that describes
how to quantize each layer weights.
"""
def __init__(
self,
ckpt_path: str,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
):
"""Initializes the loader.
Args:
ckpt_path: The filepath to the checkpoint.
is_symmetric: Whether to apply symmetric or asymmetric quantization.
attention_quant_bits: An integer that specify the target quantization bits
(support 8 or 4) for the attention layers.
feedforward_quant_bits: An integer that specify the target quantization
bits (support 8 or 4) for the feedforward layers in each Transformer
blocks.
embedding_quant_bits: An integer that specify the target quantization bits
(support 8 or 4) for the embedding (and the final projection) layers.
"""
self._ckpt_path = ckpt_path
self._is_symmetric = is_symmetric
self._attention_quant_bits = attention_quant_bits
self._feedforward_quant_bits = feedforward_quant_bits
self._embedding_quant_bits = embedding_quant_bits
def load_to_actions(
self,
) -> Iterator[Optional[List[QuantizationAction]]]:
"""Loads the checkpoint and returns the quantization actions."""
raise NotImplementedError("The load_to_actions method is not implemented.")
class LayerActionMapperBase:
"""Base class for mapping the layer weights to quantization actions.
This class is responsible for mapping from each layer to its corresponding
quantization information (e.g. target quantization bits / updated tensor
name...).
"""
def __init__(
self,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
backend: str,
):
self._is_symmetric = is_symmetric
self._attention_quant_bits = attention_quant_bits
self._feedforward_quant_bits = feedforward_quant_bits
self._embedding_quant_bits = embedding_quant_bits
self._backend = backend
def map_to_actions(
self, layer_name: str
) -> Optional[List[QuantizationAction]]:
"""Maps the layer weights to quantization actions.
Args:
layer_name: A string representing the name of the layer weight. Note that
it is expected the layer information is contained in the name which is
enough to determine the target quantization information. Any child class
is expected to implement this function.
"""
raise NotImplementedError("The map_to_actions method is not implemented.")
class ModelWriterBase:
"""Base class for writing the quantized model.
This class is responsible for taking a dictionary of the quantized
tensors/names and writing them into the format that can be loaded by the
on-device inference engine.
"""
def __init__(self, output_dir: str, backend: str):
"""Initializes the class.
Args:
output_dir: A string that represents the output directory to write the
resulting file(s).
backend: A string that represents the target backend to run the output
file(s).
"""
self._output_dir = output_dir
if not os.path.exists(self._output_dir):
os.mkdir(self._output_dir)
self._backend = backend
def write_variables(
self,
variables: Dict[str, Tuple[np.ndarray, bool]],
use_fake_values: bool = False,
):
raise NotImplementedError("The write_variables method is not implemented.")
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/converter_factory.py
+79
View File
@@ -0,0 +1,79 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Utility library that helps create the converter instances."""
from mediapipe.tasks.python.genai.converter import converter_base
from mediapipe.tasks.python.genai.converter import pytorch_converter
from mediapipe.tasks.python.genai.converter import safetensors_converter
from mediapipe.tasks.python.genai.converter import weight_bins_writer
def create_ckpt_loader(
ckpt_format: str, *args, **kwargs
) -> converter_base.CkptLoaderBase:
"""Creates the checkpoint loader.
Args:
ckpt_format: A string that indicates which input checkpoint format is.
*args: Additional arguments to be passed into the loader.
**kwargs: Additional arguments to be passed into the loader.
Returns:
A created CkptLoader instance.
"""
del args
if ckpt_format == "pytorch":
return pytorch_converter.PytorchCkptLoader(
ckpt_path=kwargs["ckpt_path"],
is_symmetric=kwargs["is_symmetric"],
attention_quant_bits=kwargs["attention_quant_bits"],
feedforward_quant_bits=kwargs["feedforward_quant_bits"],
embedding_quant_bits=kwargs["embedding_quant_bits"],
special_model=kwargs["special_model"],
backend=kwargs["backend"],
)
elif ckpt_format == "safetensors":
return safetensors_converter.SafetensorsCkptLoader(
ckpt_path=kwargs["ckpt_path"],
is_symmetric=kwargs["is_symmetric"],
attention_quant_bits=kwargs["attention_quant_bits"],
feedforward_quant_bits=kwargs["feedforward_quant_bits"],
embedding_quant_bits=kwargs["embedding_quant_bits"],
special_model=kwargs["special_model"],
backend=kwargs["backend"],
)
else:
raise ValueError(f"Unknown checkpoint format: {ckpt_format}")
def create_writer(
writer_type: str, *args, **kwargs
) -> converter_base.ModelWriterBase:
"""Creates the model writer.
Args:
writer_type: A string the indicates which model writer to create.
*args: Additional arguments to be passed into the loader.
**kwargs: Additional arguments to be passed into the loader.
Returns:
A created ModelWriter instance.
"""
del args
if writer_type == "weight_bins":
return weight_bins_writer.WeightBinsWriter(
output_dir=kwargs["output_dir"], backend=kwargs["backend"]
)
else:
raise ValueError(f"Unknown writer type: {writer_type}")
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/external_dependencies.py
+26
View File
@@ -0,0 +1,26 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
"""MediaPipe LLM Converter Dependencies."""
try:
import jax
from jax import lax
from jax import numpy as jnp
import torch
import sentencepiece
except ImportError as e:
raise ImportError(
"JAX, Torch, and SentencePiece are required for the LLM conversions."
" Please run `pip install torch jax sensetencepiece`."
) from e
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/llm_converter.py
+644
View File
@@ -0,0 +1,644 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Functions to perform the checkpoint conversion."""
import contextlib
import ctypes
import os
from typing import Any, List, Optional
from absl import logging
import numpy as np
from mediapipe.tasks.python.core import mediapipe_c_bindings
from mediapipe.tasks.python.core import mediapipe_c_utils
from mediapipe.tasks.python.core import serial_dispatcher
from mediapipe.tasks.python.genai.converter import converter_base
from mediapipe.tasks.python.genai.converter import converter_factory
from mediapipe.tasks.python.genai.converter import external_dependencies
from mediapipe.tasks.python.genai.converter import quantization_util
jnp = external_dependencies.jnp
_CTYPES_SIGNATURES = (
mediapipe_c_utils.CStatusFunction(
func_name='MpLlmConverterGenerateCpuTfLite',
core_argtypes=(
ctypes.c_char_p, # model_type
ctypes.c_char_p, # weight_path
ctypes.c_char_p, # vocab_model_file
ctypes.c_bool, # is_quantized
ctypes.c_char_p, # output_tflite_file
),
),
mediapipe_c_utils.CStatusFunction(
func_name='MpLlmConverterGenerateGpuTfLite',
core_argtypes=(
ctypes.c_char_p, # model_type
ctypes.c_char_p, # weight_path
ctypes.c_char_p, # vocab_model_file
ctypes.c_bool, # is_quantized
ctypes.c_bool, # obfuscate
ctypes.c_char_p, # output_tflite_file
ctypes.c_int, # lora_rank
ctypes.c_char_p, # lora_weight_path
ctypes.c_char_p, # lora_output_tflite_file
ctypes.c_char_p, # lora_main_model_type
ctypes.c_char_p, # image_encoder_file
ctypes.c_char_p, # image_adapter_file
ctypes.c_char_p, # submodel_type
ctypes.c_bool, # use_dynamic_ple
ctypes.c_bool, # apply_srq
ctypes.c_int, # block_size
),
),
mediapipe_c_utils.CStatusFunction(
func_name='MpLlmConverterConvertHfTokenizer',
core_argtypes=(
ctypes.c_char_p, # vocab_model_file
ctypes.c_char_p, # output_vocab_file
),
),
)
def _safe_encode_str(value: str | None) -> bytes:
"""Encodes a string to bytes, returning an empty byte string if None."""
return value.encode('utf-8') if value is not None else b''
class ConversionConfig(object):
"""Config for checkpoint conversion.
Attributes:
input_ckpt: Directory or path for the input checkpoint.
ckpt_format: Checkpoint format, e.g. 'safetensors', 'pytorch'.
model_type: Name of the model, e.g. GEMMA_2B.
backend: Target backend to run the model. Can be either "cpu" or "gpu".
output_dir: Where the output file(s) to be stored.
is_quantized: Whether the checkpoint is already quantized. If the checkpoint
is already quantized, the converter will not quantize it again and it will
ignore the quantization parameters.
is_symmetric: Whether to quantize symmetrically.
attention_quant_bits: Target quantization bits for the attention layers.
feedforward_quant_bits: Target quantization bits for the feedforward layers.
embedding_quant_bits: Target quantization bits for the embedding layers.
combine_file_only: Whether to combine the weight files only (assuming the
weight files are already existed).
vocab_model_file: The file path to the 1) SentencePiece vocab model; 2)
Hugging Face BPE tokenizer files; 1) is applicable for the Gemma model and
2) is applicable for other models. When 2) is used, the provided path is
expected to point to a directory that contains both tokenizer.json and
tokenizer_config.json files.
obfuscate: Whether to obfuscate the model.
output_tflite_file: (optional) the output tflite filename. If not provided,
the output will be `model.tflite` stored in the output_dir.
fp16_scale: A scalar value between [0, 1]. Some models can run into
activation overflow issue when running in 16-bit floating point mode. To
solve this, we need to scale down the weights of certain layers. See
go/llm-on-device-fp16 for more detailed explanation.
lora_ckpt: The directory or path for the lora checkpoint. Required in order
to convert the lora weights.
lora_rank: An integer representing the rank of LoRA. Required in order to
convert the lora weights.If not provided, then the converter assumes there
is no LoRA weights. Note that only the GPU backend supports LoRA.
lora_alpha: A float representing the scale of LoRA weights.
lora_output_tflite_file: A string indicating the name of the generated
tflite file for the LoRA weight. Only applicable when the lora_rank is not
zero.
lora_main_model_type: If the LoRA will be applied to a submodel packaged
with a main model, what the main model type is.
image_encoder_file: A string with the name of the image encoder tflite file.
image_adapter_file: A string with the name of the image adapter tflite file.
submodel_type: Name of submodel, e.g. GEMMA_2B.
use_fake_weights: Whether to use fake weights. If set to True, the weights
will be filled with zeros.
use_dynamic_ple: Whether any PLE embeddings should be loaded dynamically.
Default is true, which will cause embeddings to only be loaded into VRAM
on demand.
use_mse_quant: Whether to use MSE quantization for recomputing scales.
block_size: Default of 0 is off. Can also be 32 or 128, for using blockwise
Q4_0 compression on the models and layers which support this.
apply_srq: Default of None will apply SRQ for pre-quantized models only.
Will override this behavior if set explicitly.
"""
def __init__(
self,
input_ckpt: str,
ckpt_format: str,
model_type: str,
backend: str,
output_dir: str,
is_quantized: bool = False,
is_symmetric: bool = True,
attention_quant_bits: int = 8,
feedforward_quant_bits: int = 8,
embedding_quant_bits: int = 8,
combine_file_only: bool = False,
vocab_model_file: str = '',
obfuscate: bool = False,
output_tflite_file: Optional[str] = None,
fp16_scale: Optional[float] = None,
block_size: int = 0,
lora_ckpt: Optional[str] = None,
lora_rank: Optional[int] = None,
lora_alpha: Optional[float] = None,
lora_output_tflite_file: Optional[str] = None,
lora_main_model_type: Optional[str] = None,
image_encoder_file: Optional[str] = None,
image_adapter_file: Optional[str] = None,
submodel_type: Optional[str] = None,
use_fake_weights: bool = False,
use_dynamic_ple: bool = True,
use_mse_quant: bool = False,
apply_srq: Optional[bool] = None,
):
self.input_ckpt = input_ckpt
self.ckpt_format = ckpt_format
self.model_type = model_type
self.backend = backend
if os.path.isfile(output_dir):
raise ValueError('Output directory mush not point to an existing file.')
if not os.path.isdir(output_dir):
logging.info('Creating output directory: %s', output_dir)
os.makedirs(output_dir, exist_ok=True)
self.output_dir = output_dir
self.is_quantized = is_quantized
self.is_symmetric = is_symmetric
self.attention_quant_bits = attention_quant_bits
self.feedforward_quant_bits = feedforward_quant_bits
self.embedding_quant_bits = embedding_quant_bits
self.combine_file_only = combine_file_only
self.vocab_model_file = vocab_model_file
self.obfuscate = obfuscate
self.image_encoder_file = image_encoder_file
self.image_adapter_file = image_adapter_file
self.submodel_type = submodel_type
self.use_fake_weights = use_fake_weights
self.use_dynamic_ple = use_dynamic_ple
self.use_mse_quant = use_mse_quant
self.apply_srq = apply_srq
if output_tflite_file:
parent_dir = os.path.dirname(output_tflite_file)
if not os.path.isdir(parent_dir):
logging.info('Creating tflite parent directory: %s', parent_dir)
os.makedirs(parent_dir, exist_ok=True)
self.output_tflite_file = output_tflite_file
else:
self.output_tflite_file = os.path.join(output_dir, 'model.tflite')
self.fp16_scale = None
if block_size not in (0, 32, 128):
raise ValueError('Block size must be one of (0 (off), 32, 128)')
else:
self.block_size = block_size
self.lora_ckpt = lora_ckpt
self.lora_rank = lora_rank
self.lora_alpha = lora_alpha
self.lora_output_tflite_file = lora_output_tflite_file
self.lora_main_model_type = lora_main_model_type
if (self.lora_ckpt is None) ^ (self.lora_rank is None):
raise ValueError(
'lora_ckpt and lora_rank must be either both provided or both not'
' provided.'
)
if self.lora_rank is not None:
if backend == 'cpu':
raise ValueError('LoRA is not supported for CPU backend.')
lora_applicable_models = [
'GEMMA_2B',
'GEMMA2_2B',
'PHI_2',
'GEMMA3_1B',
'GEMMA3_4B',
'GEMMA3_12B',
'GEMMA3_27B',
'GEMMA3_300M',
]
if model_type not in lora_applicable_models:
raise ValueError(
'LoRA is only applicable for the model_type:'
f' {", ".join(lora_applicable_models)}, but get model_type:'
f' {model_type}.'
)
@contextlib.contextmanager
def filemanager(filename: str, mode: str):
try:
with open(filename, mode) as f:
yield f
finally:
pass
class _LlmConverter:
"""Bundles all conversion logic for LLM models."""
def __init__(self, lib: Any):
self._lib = serial_dispatcher.SerialDispatcher(lib, _CTYPES_SIGNATURES)
def quantize_by_actions(
self,
actions: List[converter_base.QuantizationAction],
backend: str,
is_symmetric: bool,
use_mse_quant: bool = False,
block_size: int = 0,
):
"""Quantizes the weights by actions.
Args:
actions: A list of QuantizationAction that contains the information and
tensor values to be quantized.
backend: Target backend to run the model. Can be either "cpu" or "gpu".
is_symmetric: Whether to quantize symmetrically.
use_mse_quant: Whether to use MSE quantization for recomputing scales.
block_size: 0 to disable, and 32 or 128 to use blockwise Q4_0 for
recomputing scales.
Returns:
A dictionary that maps from the updated tensor names to the quantized
tensor values + a boolean that indicates whether the tensor values need to
be packed (only applicable for the 4-bit quantized weights).
"""
output_tensors = {}
qvalue_suffix = '_quantized_value'
scale_suffix = '_quantized_scale'
zp_suffix = '_quantized_zp'
for action in actions:
if action.tensor_value is None:
continue
# The dtype needs to be compared in string as it is a custom numpy dtype.
# Explicitly cast the bfloat16 and float16 dtype to float32 to make sure
# its value is converted and serialized correctly.
if (
str(action.tensor_value.dtype) == 'bfloat16'
or action.tensor_value.dtype == np.float16
):
action.tensor_value = action.tensor_value.astype(np.float32)
if (
(not action.is_quantized)
and action.tensor_value.dtype != np.float32
and action.tensor_value.dtype != np.int8
):
raise ValueError(
'All tensors should be casted to either float32 or int8, but got: '
'%s'
% action.tensor_value.dtype
)
if action.is_quantized:
pack = action.tensor_value.dtype == jnp.int4
if qvalue_suffix in action.target_name:
target_name = action.target_name[: -len(qvalue_suffix)]
# Stores the quantized value in int8 for 4-bit quantization.
if pack:
action.tensor_value = action.tensor_value.astype(jnp.int8)
output_tensors[target_name] = (action.tensor_value, pack)
elif (
scale_suffix in action.target_name
or zp_suffix in action.target_name
):
output_tensors[action.target_name] = (
action.tensor_value,
False,
)
else:
output_tensors[action.target_name] = (action.tensor_value, False)
if action.quantize_axis:
quant_bits = action.quantize_bits
quant_axis = action.quantize_axis.copy()
if (
block_size > 0
and quant_bits == 4
and quant_axis != [0]
and quant_axis != [1]
and quant_axis != [0, 1]
):
# Turn off 4bit blockwise quant for unsupported quantization axis
# configurations
print(
'Disabling blockwise 4bit quant for unsupported quantization'
' axis: ',
quant_axis,
' on tensor ',
action.target_name,
)
quant_bits = 8
pack = quant_bits == 4
if action.tensor_value.dtype == np.int8:
if backend == 'cpu' and pack:
raise ValueError(
'Converting pre-quantized checkpoint into 4-bit is not'
' supported for CPU backend.'
)
output_tensors[action.target_name] = (action.tensor_value, pack)
else:
if is_symmetric:
bs = block_size if pack else 0
target_var, scale = quantization_util.quantize_tensor(
var=action.tensor_value,
axis=action.quantize_axis,
sym=is_symmetric,
number_bits=quant_bits,
use_mse_quant=use_mse_quant,
# blockwise quant applies to symmetric 4bit only
block_size=bs,
)
if bs > 0:
if len(scale.shape) > 3:
raise ValueError(
'Blockwise int4 quantization not yet supported for'
' 4-dimensional or higher scales.'
)
# Reshape and transpose output from blockwise quant.
# We want per-block scale values as the last dimension.
if quant_axis == [0] or quant_axis == [0, 1]:
if len(scale.shape) <= 2:
scale = scale.transpose()
else:
scale = scale.transpose([1, 2, 0])
elif quant_axis == [1]:
if len(scale.shape) == 3:
# Only tested in one case, so not sure if this generalizes.
# Seems to behave identically to our experimental
# quant_axis = [0, 1] support so far.
scale = scale.transpose([2, 0, 1])
# Reshape to fold in the per-block dimension
target_var = target_var.reshape(action.tensor_value.shape)
output_tensors[action.target_name] = (target_var, pack)
output_tensors[action.target_name + scale_suffix] = (
scale,
False,
)
zp = None
else:
target_var, scale, zp = quantization_util.quantize_tensor(
var=action.tensor_value,
axis=action.quantize_axis,
sym=is_symmetric,
number_bits=action.quantize_bits,
use_mse_quant=use_mse_quant,
)
if backend == 'cpu' and pack:
target_var, scale, zp = quantization_util.update_to_uint4(
target_var, scale, zp
)
output_tensors[action.target_name] = (target_var, pack)
output_tensors[action.target_name + scale_suffix] = (scale, False)
if zp is not None:
output_tensors[action.target_name + zp_suffix] = (zp, False)
else:
output_tensors[action.target_name] = (action.tensor_value, False)
return output_tensors
def combined_weight_bins_to_tflite(
self,
model_type: str,
backend: str,
weight_path: str,
output_tflite_file: str,
obfuscate: bool,
vocab_model_file: str,
lora_rank: Optional[int] = None,
lora_weight_path: Optional[str] = None,
lora_output_tflite_file: Optional[str] = None,
lora_main_model_type: Optional[str] = None,
image_encoder_file: Optional[str] = None,
image_adapter_file: Optional[str] = None,
submodel_type: Optional[str] = None,
use_dynamic_ple: Optional[bool] = None,
apply_srq: Optional[bool] = None,
block_size: int = 0,
):
"""Combines weight files to tflite file."""
if backend == 'cpu':
if lora_rank is not None:
logging.fatal('LoRA is not supported for CPU backend.')
self._lib.MpLlmConverterGenerateCpuTfLite(
model_type.encode('utf-8'),
weight_path.encode('utf-8'),
vocab_model_file.encode('utf-8'),
True,
output_tflite_file.encode('utf-8'),
)
elif backend == 'gpu':
self._lib.MpLlmConverterGenerateGpuTfLite(
model_type.encode('utf-8'),
weight_path.encode('utf-8'),
vocab_model_file.encode('utf-8'),
True,
obfuscate,
output_tflite_file.encode('utf-8'),
0 if lora_rank is None else lora_rank,
_safe_encode_str(lora_weight_path),
_safe_encode_str(lora_output_tflite_file),
_safe_encode_str(lora_main_model_type),
_safe_encode_str(image_encoder_file),
_safe_encode_str(image_adapter_file),
_safe_encode_str(submodel_type),
True if use_dynamic_ple is None else use_dynamic_ple,
False if apply_srq is None else apply_srq,
block_size,
)
else:
raise ValueError(f'Unsupported backend: {backend}')
def convert_bpe_vocab(self, vocab_model_file: str, output_dir: str) -> str:
"""Converts the BPE vocab model file to SPM format.
Args:
vocab_model_file: The input BPE vocab model file path.
output_dir: The output directory to store the SPM vocab model file.
Returns:
The path to the output SPM vocab model file.
"""
if not os.path.isdir(vocab_model_file):
raise ValueError(
'The input BPE vocab model file path is expected to be a directory'
' that contains both tokenizer.json and tokenizer_config.json files.'
)
output_vocab_file = os.path.join(output_dir, 'spm.model')
self._lib.MpLlmConverterConvertHfTokenizer(
vocab_model_file.encode('utf-8'),
output_vocab_file.encode('utf-8'),
)
return output_vocab_file
def sort_layer_info(self, layer_info_file: str) -> None:
"""Loads and sorts the layer info file."""
layer_info = []
with filemanager(layer_info_file, 'r') as finfo:
for line in finfo:
line = line.strip()
if line:
layer_info.append(line)
layer_info = list(set(layer_info))
layer_info.sort()
with filemanager(layer_info_file, 'w') as finfo:
for line in layer_info:
finfo.write(line + '\n')
finfo.write('\n')
def maybe_quantize_and_write_tensors_to_bins(
self,
ckpt_loader: converter_base.CkptLoaderBase,
config: ConversionConfig,
) -> None:
"""Quantizes the weight tensors according to the loader and writes them to bins."""
actions = ckpt_loader.load_to_actions()
for action in actions:
# Quantize the weight
quantized_tensors = self.quantize_by_actions(
action,
config.backend,
config.is_symmetric,
config.use_mse_quant,
block_size=config.block_size,
)
del action
# Write the tensors into file(s).
writer = converter_factory.create_writer(
writer_type='weight_bins',
output_dir=config.output_dir,
backend=config.backend,
)
writer.write_variables(quantized_tensors, config.use_fake_weights)
del quantized_tensors
del writer
def convert_checkpoint(self, config: ConversionConfig) -> None:
"""Converts the checkpoint to tflite file."""
logging.info('input folder: %s', config.input_ckpt)
if os.path.isdir(config.vocab_model_file):
vocab_model_path = self.convert_bpe_vocab(
config.vocab_model_file, config.output_dir
)
else:
vocab_model_path = config.vocab_model_file
if not config.combine_file_only:
# Load the layer weights and prepare the quantization configurations.
loader = converter_factory.create_ckpt_loader(
config.ckpt_format,
ckpt_path=config.input_ckpt,
is_quantized=config.is_quantized,
is_symmetric=config.is_symmetric,
backend=config.backend,
attention_quant_bits=config.attention_quant_bits,
feedforward_quant_bits=config.feedforward_quant_bits,
embedding_quant_bits=config.embedding_quant_bits,
special_model=config.model_type,
fp16_scale=config.fp16_scale,
)
self.maybe_quantize_and_write_tensors_to_bins(loader, config)
if config.lora_ckpt is not None and config.lora_ckpt != config.input_ckpt:
# If lora ckpt and the input ckpt is the same. The lora conversion is
# handled in the previous loader.
lora_loader = converter_factory.create_ckpt_loader(
config.ckpt_format,
ckpt_path=config.lora_ckpt,
is_quantized=config.is_quantized,
is_symmetric=config.is_symmetric,
backend=config.backend,
attention_quant_bits=config.attention_quant_bits,
feedforward_quant_bits=config.feedforward_quant_bits,
embedding_quant_bits=config.embedding_quant_bits,
special_model=config.model_type,
)
self.maybe_quantize_and_write_tensors_to_bins(lora_loader, config)
self.sort_layer_info(os.path.join(config.output_dir, 'layer_info.txt'))
self.combined_weight_bins_to_tflite(
config.model_type,
config.backend,
weight_path=config.output_dir,
output_tflite_file=config.output_tflite_file,
obfuscate=config.obfuscate,
vocab_model_file=vocab_model_path,
lora_rank=config.lora_rank,
lora_weight_path=config.output_dir,
lora_output_tflite_file=config.lora_output_tflite_file,
lora_main_model_type=config.lora_main_model_type,
image_encoder_file=config.image_encoder_file,
image_adapter_file=config.image_adapter_file,
submodel_type=config.submodel_type,
use_dynamic_ple=config.use_dynamic_ple,
# Default SRQ behavior is to match is_quantized
apply_srq=config.is_quantized
if config.apply_srq is None
else config.apply_srq,
block_size=config.block_size,
)
def __enter__(self) -> '_LlmConverter':
"""Returns `self` upon entering the runtime context."""
return self
def __exit__(self, *args) -> None:
"""Shuts down the LlmConverter on exit of the context manager.
Args:
*args: Unused.
Raises:
RuntimeError: If the LLM converter failed to close.
"""
del args # Unused.
self._lib.close()
def convert_checkpoint(config: ConversionConfig) -> None:
"""Converts the checkpoint to tflite file."""
lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES)
with _LlmConverter(lib) as converter:
converter.convert_checkpoint(config)
def quantize_by_actions(
actions: List[converter_base.QuantizationAction],
backend: str,
is_symmetric: bool,
use_mse_quant: bool = False,
block_size: int = 0,
):
"""Quantizes the weights by actions.
Args:
actions: A list of QuantizationAction that contains the information and
tensor values to be quantized.
backend: Target backend to run the model. Can be either "cpu" or "gpu".
is_symmetric: Whether to quantize symmetrically.
use_mse_quant: Whether to use MSE quantization for recomputing scales.
block_size: 0 to disable, and 32 or 128 to use blockwise quantization for
recomputing scales.
Returns:
A dictionary that maps from the updated tensor names to the quantized
tensor values + a boolean that indicates whether the tensor values need to
be packed (only applicable for the 4-bit quantized weights).
"""
lib = mediapipe_c_bindings.load_raw_library(_CTYPES_SIGNATURES)
with _LlmConverter(lib) as converter:
return converter.quantize_by_actions(
actions, backend, is_symmetric, use_mse_quant, block_size
)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/llm_converter_test.py
+264
View File
@@ -0,0 +1,264 @@
# Copyright 2025 The MediaPipe Authors.
#
# 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.
from unittest import mock
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
import unittest
from mediapipe.tasks.python.genai.converter import converter_base
from mediapipe.tasks.python.genai.converter import llm_converter
_LLM_CONVERTER_FUNCTIONS = (
'MpLlmConverterGenerateCpuTfLite',
'MpLlmConverterGenerateGpuTfLite',
'MpLlmConverterConvertHfTokenizer',
)
class LlmConverterTest(googletest.TestCase, parameterized.TestCase):
def get_fake_action(self, input_dtype):
if input_dtype == 'bfloat16':
# Create a TensorFlow bfloat16 tensor
bfloat16_tensor = tf.constant([1.0, -1.0, 2.0, -2.0], dtype=tf.bfloat16)
# Convert the TensorFlow tensor to a NumPy array
tensor_value = bfloat16_tensor.numpy()
else:
tensor_value = np.array(
[1.0, -1.0, 2.0, -2.0], dtype=np.dtype(input_dtype)
)
return converter_base.QuantizationAction(
tensor_name='mdl_vars.params.lm.softmax.logits_ffn.w',
target_name='params.lm.softmax.logits_ffn.w',
quantize_axis=[0],
quantize_bits=8,
pack_dim=0,
tensor_value=tensor_value,
)
@parameterized.parameters(
{'input_dtype': 'float32'},
{'input_dtype': 'float16'},
{'input_dtype': 'bfloat16'},
{'input_dtype': 'int8'},
)
def test_quantize_by_actions(self, input_dtype):
out = llm_converter.quantize_by_actions(
[self.get_fake_action(input_dtype)], backend='gpu', is_symmetric=True
)
if input_dtype == 'int8':
# The values are pre-quantized and should be the same.
np.testing.assert_allclose(
out['params.lm.softmax.logits_ffn.w'][0],
np.array([1, -1, 2, -2], dtype=np.int8),
)
else:
np.testing.assert_allclose(
out['params.lm.softmax.logits_ffn.w'][0],
np.array([64, -64, 127, -127], dtype=np.int8),
)
np.testing.assert_allclose(
out['params.lm.softmax.logits_ffn.w_quantized_scale'][0],
np.array(0.015748, dtype=np.float32),
rtol=1e-03,
)
def test_generate_cpu_tflite_receives_correct_args(self):
c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
c_lib.MpLlmConverterGenerateCpuTfLite.return_value = 0
config = llm_converter.ConversionConfig(
model_type='GEMMA_2B',
backend='cpu',
input_ckpt='/tmp/ckpt',
ckpt_format='fake',
output_dir='/tmp/',
vocab_model_file='/tmp/vocab.model',
output_tflite_file='/tmp/output.tflite',
)
llm_converter_lib = llm_converter._LlmConverter(c_lib)
llm_converter_lib.combined_weight_bins_to_tflite(
model_type=config.model_type,
backend=config.backend,
weight_path=config.output_dir,
output_tflite_file=config.output_tflite_file,
obfuscate=config.obfuscate,
vocab_model_file=config.vocab_model_file,
)
c_lib.MpLlmConverterGenerateCpuTfLite.assert_called_once_with(
b'GEMMA_2B',
b'/tmp/',
b'/tmp/vocab.model',
True,
b'/tmp/output.tflite',
mock.ANY, # error_message
)
def test_generate_cpu_tflite_propagates_failure(self):
c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
c_lib.MpLlmConverterGenerateCpuTfLite.return_value = 13 # Simulate failure
config = llm_converter.ConversionConfig(
model_type='GEMMA_2B',
backend='cpu',
input_ckpt='/tmp/ckpt',
ckpt_format='fake',
output_dir='/tmp/',
vocab_model_file='/tmp/vocab.model',
output_tflite_file='/tmp/output.tflite',
)
llm_converter_lib = llm_converter._LlmConverter(c_lib)
with self.assertRaises(RuntimeError):
llm_converter_lib.combined_weight_bins_to_tflite(
model_type=config.model_type,
backend=config.backend,
weight_path=config.output_dir,
output_tflite_file=config.output_tflite_file,
obfuscate=config.obfuscate,
vocab_model_file=config.vocab_model_file,
)
def test_generate_gpu_tflite_receives_correct_args(self):
c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
c_lib.MpLlmConverterGenerateGpuTfLite.return_value = 0
config = llm_converter.ConversionConfig(
model_type='GEMMA_2B',
backend='gpu',
input_ckpt='/tmp/ckpt',
ckpt_format='fake',
output_dir='/tmp/',
vocab_model_file='/tmp/vocab.model',
output_tflite_file='/tmp/output.tflite',
use_dynamic_ple=True,
)
llm_converter_lib = llm_converter._LlmConverter(c_lib)
llm_converter_lib.combined_weight_bins_to_tflite(
model_type=config.model_type,
backend=config.backend,
weight_path=config.output_dir,
output_tflite_file=config.output_tflite_file,
obfuscate=config.obfuscate,
vocab_model_file=config.vocab_model_file,
lora_rank=4,
lora_weight_path='/tmp/lora_weights',
lora_output_tflite_file='/tmp/lora_output.tflite',
lora_main_model_type='GEMMA_2B',
image_encoder_file='/tmp/image_encoder.tflite',
image_adapter_file='/tmp/image_adapter.tflite',
submodel_type='GEMMA_2B',
use_dynamic_ple=config.use_dynamic_ple,
apply_srq=False,
)
c_lib.MpLlmConverterGenerateGpuTfLite.assert_called_once_with(
b'GEMMA_2B',
b'/tmp/',
b'/tmp/vocab.model',
True,
False,
b'/tmp/output.tflite',
4,
b'/tmp/lora_weights',
b'/tmp/lora_output.tflite',
b'GEMMA_2B',
b'/tmp/image_encoder.tflite',
b'/tmp/image_adapter.tflite',
b'GEMMA_2B',
True,
False,
0,
mock.ANY, # error_message
)
def test_generate_gpu_tflite_propagates_failure(self):
c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
c_lib.MpLlmConverterGenerateGpuTfLite.return_value = 13 # Simulate failure
config = llm_converter.ConversionConfig(
model_type='GEMMA_2B',
backend='gpu',
input_ckpt='/tmp/ckpt',
ckpt_format='fake',
output_dir='/tmp/',
vocab_model_file='/tmp/vocab.model',
output_tflite_file='/tmp/output.tflite',
use_dynamic_ple=True,
)
llm_converter_lib = llm_converter._LlmConverter(c_lib)
with self.assertRaises(RuntimeError):
llm_converter_lib.combined_weight_bins_to_tflite(
model_type=config.model_type,
backend=config.backend,
weight_path=config.output_dir,
output_tflite_file=config.output_tflite_file,
obfuscate=config.obfuscate,
vocab_model_file=config.vocab_model_file,
lora_rank=4,
lora_weight_path='/tmp/lora_weights',
lora_output_tflite_file='/tmp/lora_output.tflite',
lora_main_model_type='GEMMA_2B',
image_encoder_file='/tmp/image_encoder.tflite',
image_adapter_file='/tmp/image_adapter.tflite',
submodel_type='GEMMA_2B',
use_dynamic_ple=config.use_dynamic_ple,
apply_srq=False,
)
@mock.patch('os.path.isdir', return_value=True)
@mock.patch('os.path.join', return_value='/tmp/spm.model')
def test_convert_hf_tokenizer_receives_correct_args(
self, mock_join, mock_isdir
):
del mock_join, mock_isdir
c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
c_lib.MpLlmConverterConvertHfTokenizer.return_value = 0
llm_converter_lib = llm_converter._LlmConverter(c_lib)
output_file = llm_converter_lib.convert_bpe_vocab(
vocab_model_file='/tmp/hf_tokenizer', output_dir='/tmp'
)
c_lib.MpLlmConverterConvertHfTokenizer.assert_called_once_with(
b'/tmp/hf_tokenizer', b'/tmp/spm.model', mock.ANY
)
self.assertEqual(output_file, '/tmp/spm.model')
@mock.patch('os.path.isdir', return_value=True)
@mock.patch('os.path.join', return_value='/tmp/spm.model')
def test_convert_hf_tokenizer_propagates_failure(self, mock_join, mock_isdir):
del mock_join, mock_isdir
c_lib = mock.MagicMock(spec_set=_LLM_CONVERTER_FUNCTIONS)
c_lib.MpLlmConverterConvertHfTokenizer.return_value = 13 # Simulate failure
llm_converter_lib = llm_converter._LlmConverter(c_lib)
with self.assertRaises(RuntimeError):
llm_converter_lib.convert_bpe_vocab(
vocab_model_file='/tmp/hf_tokenizer', output_dir='/tmp'
)
if __name__ == '__main__':
googletest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/pytorch_converter.py
+320
View File
@@ -0,0 +1,320 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""CkptLoader implementation for loading the Pytorch file."""
from typing import Iterator
import enum
import os
from typing import List, Optional
import numpy as np
from mediapipe.tasks.python.genai.converter import converter_base
from mediapipe.tasks.python.genai.converter import external_dependencies
torch = external_dependencies.torch
class _PytorchReader:
"""Pytorch reader."""
def __init__(self, model_path: str):
if not os.path.exists(model_path):
raise ValueError(f"{model_path} does not exists.")
self._model = torch.load(model_path, map_location=torch.device("cpu"))
def read_tensor_as_numpy(self, tensor_name) -> np.ndarray:
tensor = (
self._model[tensor_name]
.to(torch.float32)
.t()
.contiguous()
.detach()
.cpu()
.numpy()
)
return tensor
def get_tensor_names(self) -> List[str]:
names = list(self._model.keys())
return names
class LayerType(enum.Enum):
"""Enum for layer type."""
NONE = 0
ATTENTION = 1 # Layer is part of the attention module.
FEEDFORWARD = 2 # Layer is part of the feedforward module in the Transformer.
EMBEDDING = 3 # Layer is the embedding lookup or final projection layer.
LAYER_NORM = (
4 # Layer is layer normalization before and after attention layer.
)
@classmethod
def get_layer_type(cls, layer_name: str):
"""Gets the layer type of the given layer name."""
ffn_layers = [
"mlp",
]
attn_layers = [
"self_attention",
]
emb_layers = [
"word_embeddings",
"lm_head",
]
layer_norms = [
"input_layernorm",
"post_attention_layernorm",
"ln_f",
]
if any(sub_name in layer_name for sub_name in attn_layers):
return LayerType.ATTENTION
if any(sub_name in layer_name for sub_name in ffn_layers):
return LayerType.FEEDFORWARD
if any(sub_name in layer_name for sub_name in emb_layers):
return LayerType.EMBEDDING
if any(sub_name in layer_name for sub_name in layer_norms):
return LayerType.LAYER_NORM
else:
return LayerType.NONE
class FalconMapper(converter_base.LayerActionMapperBase):
"""LayerActionMapper for handling the Falcon-rw-1b model."""
def __init__(
self,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
backend: str,
reader: _PytorchReader,
):
super().__init__(
is_symmetric=is_symmetric,
attention_quant_bits=attention_quant_bits,
feedforward_quant_bits=feedforward_quant_bits,
embedding_quant_bits=embedding_quant_bits,
backend=backend,
)
self._reader = reader
def map_to_actions(
self, layer_name: str
) -> Optional[List[converter_base.QuantizationAction]]:
"""Map the given layer name to actions."""
actions = []
tensor_value = self._reader.read_tensor_as_numpy(layer_name)
if "query_key_value" in layer_name:
qkv_tensors = self._decompose_falcon_qkv(tensor_value)
for tensor, name in zip(qkv_tensors, ["q", "k", "v"]):
decomposed_name = layer_name.replace("query_key_value", name)
action = self._map_to_action_helper(tensor, decomposed_name)
actions.append(action)
else:
actions.append(self._map_to_action_helper(tensor_value, layer_name))
return actions
def _map_to_action_helper(
self, tensor_value: np.ndarray, layer_name: str
) -> converter_base.QuantizationAction:
quantize_axis = None
quantize_bits = None
layer_type = LayerType.get_layer_type(layer_name)
if layer_type != LayerType.LAYER_NORM and layer_name.endswith(".weight"):
layer_type = LayerType.get_layer_type(layer_name)
quantize_axis = [0]
if layer_type == LayerType.FEEDFORWARD:
quantize_bits = self._feedforward_quant_bits
elif layer_type == LayerType.ATTENTION:
quantize_bits = self._attention_quant_bits
if self._backend == "cpu" and ".dense." in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
elif layer_type == LayerType.EMBEDDING:
quantize_bits = self._embedding_quant_bits
if self._backend == "cpu" and "word_embeddings" in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
target_name = self.update_target_name(layer_name)
return converter_base.QuantizationAction(
tensor_name=layer_name,
tensor_value=tensor_value,
target_name=target_name,
quantize_axis=quantize_axis,
quantize_bits=quantize_bits,
pack_dim=0,
)
def _decompose_falcon_qkv(self, tensor_value: np.ndarray) -> List[np.ndarray]:
"""Decomposes combined qkv tensor used in falcon model into separate q, k and v tensors."""
chunk_size = 64
hidden_size = 2048
tensor_value = tensor_value.transpose()
q_tensor = np.zeros(
(hidden_size,)
+ ((hidden_size,) if len(tensor_value.shape) == 2 else ()),
dtype=tensor_value.dtype,
)
k_tensor = np.zeros_like(q_tensor, dtype=tensor_value.dtype)
v_tensor = np.zeros_like(k_tensor, dtype=tensor_value.dtype)
j = 0
for i in range(0 * chunk_size, hidden_size * 3, chunk_size * 3):
q_tensor[j : j + chunk_size] = tensor_value[i : i + chunk_size]
j += chunk_size
j = 0
for i in range(1 * chunk_size, hidden_size * 3, chunk_size * 3):
k_tensor[j : j + chunk_size] = tensor_value[i : i + chunk_size]
j += chunk_size
j = 0
for i in range(2 * chunk_size, hidden_size * 3, chunk_size * 3):
v_tensor[j : j + chunk_size] = tensor_value[i : i + chunk_size]
j += chunk_size
return [
np.ascontiguousarray(q_tensor.transpose()),
np.ascontiguousarray(k_tensor.transpose()),
np.ascontiguousarray(v_tensor.transpose()),
]
def update_target_name(self, target_name: str) -> str:
"""Updates the target name to match the tensor name convention."""
layer_type = LayerType.get_layer_type(target_name)
target_name = target_name.replace(
"transformer.h.", "params.lm.transformer.x_layers_"
)
if layer_type == LayerType.FEEDFORWARD:
target_name = target_name.replace(".weight", ".linear.w")
target_name = target_name.replace(".bias", ".bias.b")
target_name = target_name.replace(
"mlp.dense_h_to_4h", "ff_layer.ffn_layer1"
)
target_name = target_name.replace(
"mlp.dense_4h_to_h", "ff_layer.ffn_layer2"
)
elif layer_type == LayerType.ATTENTION:
target_name = target_name.replace("dense", "post")
target_name = target_name.replace(".weight", ".linear.w")
target_name = target_name.replace(".bias", ".bias.b")
elif layer_type == LayerType.EMBEDDING:
target_name = target_name.replace(
"transformer.word_embeddings", "params.lm.token_embedding"
)
target_name = target_name.replace(
"lm_head", "params.lm.softmax.logits_ffn"
)
target_name = target_name.replace(".weight", ".w")
elif layer_type == LayerType.LAYER_NORM:
target_name = target_name.replace("input_layernorm", "pre_layer_norm")
target_name = target_name.replace(
"pre_layer_norm.weight", "pre_layer_norm.scale"
)
if self._backend == "cpu":
target_name = target_name.replace(
"post_attention_layernorm", "ff_layer.pre_layer_norm"
)
target_name = target_name.replace(
"ff_layer.pre_layer_norm.weight", "ff_layer.pre_layer_norm.scale"
)
else:
target_name = target_name.replace(
"post_attention_layernorm", "post_layer_norm"
)
target_name = target_name.replace(
"post_layer_norm.weight", "post_layer_norm.scale"
)
target_name = target_name.replace(
"transformer.ln_f.weight", "params.lm.final_ln.scale"
)
target_name = target_name.replace(
"transformer.ln_f.bias", "params.lm.final_ln.bias"
)
return target_name
class PytorchCkptLoader(converter_base.CkptLoaderBase):
"""CkptLoader implementation for loading the Pytorch model."""
def __init__(
self,
ckpt_path: str,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
special_model: str,
backend: str,
):
"""Initializes the loader.
Args:
ckpt_path: The filepath to the safetensors file.
is_symmetric: Whether to apply symmetric or asymmetric quantization.
attention_quant_bits: An integer that specify the target quantization bits
(support 8 or 4) for the attention layers.
feedforward_quant_bits: An integer that specify the target quantization
bits (support 8 or 4) for the feedforward layers in each Transformer
blocks.
embedding_quant_bits: An integer that specify the target quantization bits
(support 8 or 4) for the embedding (and the final projection) layers.
special_model: A string that indicates which input model is and whether
any special treatment is needed.
backend: A string indicating the backend used when converting this model.
Valid options are "cpu" and "gpu".
"""
super().__init__(
ckpt_path,
is_symmetric,
attention_quant_bits,
feedforward_quant_bits,
embedding_quant_bits,
)
self._special_model = special_model
self._reader = _PytorchReader(ckpt_path)
if special_model in ["FALCON_RW_1B"]:
self.mapper = FalconMapper(
is_symmetric,
attention_quant_bits,
feedforward_quant_bits,
embedding_quant_bits,
backend,
self._reader,
)
else:
raise ValueError(f"Unknown special model: {special_model}")
def load_to_actions(
self,
) -> Iterator[List[converter_base.QuantizationAction]]:
tensor_names = self._reader.get_tensor_names()
for tensor_name in tensor_names:
tensor_actions = self.mapper.map_to_actions(tensor_name)
if tensor_actions is None:
continue
yield tensor_actions
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/pytorch_converter_test.py
+86
View File
@@ -0,0 +1,86 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Unit tests for pytorch_converter."""
import os
from absl.testing import absltest
from absl.testing import parameterized
from mediapipe.tasks.python.genai.converter import pytorch_converter
from mediapipe.tasks.python.test import test_utils
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
_PYTORCH_FILE = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, 'falcon_rw_1b_test_weight.pt')
)
class PytorchConverterTest(parameterized.TestCase):
VARIABLE_NAMES = [
'transformer.word_embeddings.weight',
'transformer.h.0.input_layernorm.weight',
'transformer.h.0.input_layernorm.bias',
'transformer.h.0.self_attention.query_key_value.weight',
'transformer.h.0.self_attention.query_key_value.bias',
'transformer.h.0.self_attention.dense.weight',
'transformer.h.0.self_attention.dense.bias',
'transformer.h.0.post_attention_layernorm.weight',
'transformer.h.0.post_attention_layernorm.bias',
'transformer.h.0.mlp.dense_h_to_4h.weight',
'transformer.h.0.mlp.dense_h_to_4h.bias',
'transformer.h.0.mlp.dense_4h_to_h.weight',
'transformer.h.0.mlp.dense_4h_to_h.bias',
'transformer.ln_f.weight',
'transformer.ln_f.bias',
'lm_head.weight',
]
def test_init(self):
loader = pytorch_converter.PytorchCkptLoader(
ckpt_path=_PYTORCH_FILE,
is_symmetric=True,
attention_quant_bits=8,
feedforward_quant_bits=8,
embedding_quant_bits=8,
special_model='FALCON_RW_1B',
backend='cpu',
)
self.assertEqual(loader._ckpt_path, _PYTORCH_FILE)
self.assertEqual(loader._is_symmetric, True)
self.assertEqual(loader._attention_quant_bits, 8)
self.assertEqual(loader._feedforward_quant_bits, 8)
@parameterized.product(
quant_bits=(4, 8),
)
def test_load_to_actions(self, quant_bits):
loader = pytorch_converter.PytorchCkptLoader(
ckpt_path=_PYTORCH_FILE,
is_symmetric=True,
attention_quant_bits=8,
feedforward_quant_bits=quant_bits,
embedding_quant_bits=8,
special_model='FALCON_RW_1B',
backend='cpu',
)
actions = loader.load_to_actions()
# There are 16 layers in the model, but qkv weight and bias would be
# decomposed to q, k, v tensors, so there would be 20 quantization actions.
self.assertEqual(sum(len(action) for action in actions), 20)
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/quantization_util.py
+694
View File
@@ -0,0 +1,694 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Utilities for quantizing tensors.
Note that this is a reduced fork version of the praxis libraries to provide a
self-contained library for packaging.
"""
from typing import Any, List, Optional, Sequence, Tuple, Union
import numpy as np
from mediapipe.tasks.python.genai.converter import external_dependencies
jax = external_dependencies.jax
jnp = external_dependencies.jnp
lax = external_dependencies.lax
JTensor = jax.Array
_UINT4_ZP = 8 # Default zero point for unsigned 4-bit.
# The following values are taken from fake_quantizer.py in Gemax.
_MSE_QUANT_MULS = {
8: 0.05408,
4: 0.37755,
}
def _get_scan_range() -> np.ndarray:
# Produce candidate scan values.
return np.linspace(1.0, 0.5, num=11)
def _get_mean_error(bound, t, min_value, max_value, p_value):
scale = bound / max_value
candidate = jnp.divide(t, scale)
candidate = jnp.clip(jnp.round(candidate), min_value, max_value)
candidate = jnp.multiply(candidate, scale)
pmean_error = jnp.mean(jnp.abs(jnp.subtract(candidate, t)) ** p_value)
return pmean_error
def _get_best_bound_per_tensor(
t: JTensor,
bound: JTensor,
min_value: float,
max_value: float,
p_value: float = 1.0,
) -> JTensor:
"""Scan around [0.5, 1] * hard max value to get bound value for whole tensor.
This does a scan to get bound value(s) that minimize mean absolute error (MAE)
between original tensor 't' and quantized tensor. It's (almost) equivalent to
maximizing entropy.
Args:
t: The input float tensor.
bound: The hard max value for tensor 't'. It has the same length as shape.
min_value: Minimal value for the quantization bound.
max_value: Maximal value for the quantization bound.
p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
Returns:
The best bound values for 't', that minimize p-mean error.
"""
def _quant(scaling_factors):
return _get_mean_error(
bound * scaling_factors, t, min_value, max_value, p_value
)
scaling_factors = _get_scan_range()
diffs = jax.vmap(_quant)(scaling_factors)
best_scaling = scaling_factors[jnp.argmin(diffs)].astype(bound.dtype)
return bound * best_scaling
def _quantrow(
vec: JTensor,
bound: JTensor,
min_value: float,
max_value: float,
p_value: float,
factors: np.ndarray,
) -> JTensor:
"""Get best rescaling factor from a list of factors applied a channel.
Args:
vec: The vector in a channel.
bound: The hard bound (max(abs(vec))) of the vector.
min_value: The target min value.
max_value: The target max value.
p_value: Exponent of the p-mean error metric.
factors: The values to be applied on top of bound.
Returns:
adjusted bound value out of the list of factors applied to bound.
"""
def _quant(bounds):
return _get_mean_error(bounds, vec, min_value, max_value, p_value)
diffs = jax.vmap(_quant)(bound * factors)
best_scaling = factors[jnp.argmin(diffs)]
return bound * best_scaling
def _get_best_bound_per_channel(
t: JTensor,
bound: JTensor,
min_value: float,
max_value: float,
p_value: float = 1.0,
) -> JTensor:
"""Scan around [0.5, 1] * hard max value to get bound value for each channel.
This does a scan to get bound value(s) that minimize mean absolute error (MAE)
between original tensor 't' and quantized tensor. It's (almost) equivalent to
maximizing entropy.
Args:
t: The input float tensor.
bound: The hard max value for tensor 't'. It has the same length as shape.
min_value: Minimal value for the quantization bound.
max_value: Maximal value for the quantization bound.
p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
Returns:
The best bound values for 't', that minimize p-mean error.
"""
assert len(t.shape) == 2
assert len(bound.shape) == 2
assert t.shape[1] == bound.shape[1]
assert bound.shape[0] == 1
scans = _get_scan_range()
def _quant(tensor, bound, min_value, max_value, p_value, factors):
ret = np.zeros(bound.shape)
for i in range(len(tensor)):
best = _quantrow(
tensor[i], bound[i], min_value, max_value, p_value, factors
)
ret[i] = best
return ret
t = t.transpose()
t_split = list(t)
res = _quant(t_split, bound[0, :], min_value, max_value, p_value, scans)
res = res.reshape(bound.shape)
return res
def get_best_bound(
t: JTensor,
bound: JTensor,
min_value: float,
max_value: float,
p_value: float = 1.0,
per_channel: bool = False,
) -> JTensor:
"""Scan multiple factors on max value to get best bound value.
This does a scan to get bound value(s) that minimize mean absolute error (MAE)
between original tensor 't' and quantized tensor. It's (almost) equivalent to
maximizing entropy.
Args:
t: The input float tensor.
bound: The hard max value for tensor 't'. It has the same length as shape.
min_value: Minimal value for the quantization bound.
max_value: Maximal value for the quantization bound.
p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
per_channel: if get best bound for entire tensor or per channel.
Returns:
The best bound values for 't', that minimize p-mean error.
"""
if per_channel:
return _get_best_bound_per_channel(t, bound, min_value, max_value, p_value)
else:
return _get_best_bound_per_tensor(t, bound, min_value, max_value, p_value)
def get_min_max(
bits: int = 8,
unsigned: bool = False,
use_fp: bool = False,
) -> Tuple[float, float]:
"""Gets the min/max range for a given number of bits.
Args:
bits: Target number of bits for quantization.
unsigned: If True compute min and max for unsigned number, else for signed.
use_fp: in floating point.
Returns:
min/max values for the provide number of bits.
"""
if use_fp:
# TODO: support other fp types.
return -448.0, 448.0
# Calculation instead of jax.iinfo is used to support bits beside 4 and 8.
if unsigned:
# For unsigned 8 bits precision it is [0, 255]
return 0, 2**bits - 1
else:
# For signed 8 bits precision it is [-128, 127]
return -1 * 2 ** (bits - 1), 2 ** (bits - 1) - 1
def pass_through(x: JTensor, fn: Any) -> JTensor:
# Create an exactly-zero expression with Sterbenz lemma that has an
# exactly-one gradient.
return x - jax.lax.stop_gradient(x) + jax.lax.stop_gradient(fn(x))
def mse_reduce_precision(
t: JTensor,
contract_dims: Optional[Sequence[int]],
need_gradient: bool = False,
bits: int = 8,
use_symmetric: bool = True,
use_fp: bool = False,
) -> Tuple[JTensor, JTensor, Optional[JTensor]]:
"""Reduce the precision of a tensor using MSE quantization.
Args:
t: Input tensor.
contract_dims: Specifies contracting dimensions of the input tensor.
need_gradient: If gradient is needed out of this function.
bits: Target number of bits.
use_symmetric: If the input tensor is quantized symmetrically.
use_fp: Use floating point.
Returns:
A tuple of quantized tensor, quantization scale
and quantization zero point (optional).
"""
if need_gradient:
raise ValueError('need_gradient is not supported for mse_reduce_precision.')
if use_fp:
raise ValueError('use_fp is not supported for mse_reduce_precision.')
if not use_symmetric:
raise ValueError('Gemax mse_quant is typically symmetric.')
scale_multiplier = _MSE_QUANT_MULS.get(bits)
if scale_multiplier is None:
raise ValueError(f'MSE_QUANT_MULS not defined for {bits} bits.')
rms = jnp.sqrt(jnp.mean(t**2, axis=contract_dims, keepdims=True))
scale = scale_multiplier * rms + 1e-12 # Add epsilon as in Gemax
# Calculate the zero point used in fake_quantizer.mse_quant with narrow range.
maxq = 2**bits - 1
zero_mse_quant = (maxq - 1) / 2.0 # This is 7 for 4 bits
# Simulate the quantization steps from fake_quantizer.mse_quant:
# 1. Scale and shift: t / scale + zero_mse_quant
t_shifted = jnp.divide(t, scale) + zero_mse_quant
# 2. Rounding
q_rounded = jnp.round(t_shifted)
# 3. Clipping to [0, maxq - 1]
q_clipped = jnp.clip(q_rounded, 0, maxq - 1) # [0, 14] for 4 bits
# The quantized value (qvar) to be stored should be
# `q_clipped - zero_mse_quant`. This makes the dequantization match:
# (qvar - zp) * scale == (q_clipped - zero_mse_quant) * scale
qvar = q_clipped - zero_mse_quant # Ranges from [-7, 7] for 4 bits
# ml_drift expects a zero point for dequantization. To match the effective
# dequantization of mse_quant, we set the zero point to 0 (or None).
zp = None
container_dtype = (
jnp.int8 if bits <= 8 else jnp.int16 if bits <= 16 else jnp.int32
)
qvar = qvar.astype(container_dtype)
return qvar, scale, zp
def reduce_precision(
t: JTensor,
contract_dims: Optional[Sequence[int]],
need_gradient: bool = False,
bits: int = 8,
optimization_on_bound: bool = False,
p_value: float = 1.0,
percentile: float = 1.0,
use_symmetric: bool = True,
use_fp: bool = False,
add_scale_eps: bool = False,
per_channel: bool = False,
random_rounding: bool = False,
key: Optional[jax.Array] = None,
) -> Tuple[JTensor, JTensor, Optional[JTensor]]:
"""Reduce the precision of a tensor.
Generic for all tensors.
Args:
t: Input tensor.
contract_dims: Specifies contracting dimensions of the input tensor.
need_gradient: If gradient is needed out of this function.
bits: Target number of bits.
optimization_on_bound: If MAE bound optimizer is used.
p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
percentile: Percentile Factor to apply on the min/max range. Setting this to
other than 1.0 disables optimization_on_bound.
use_symmetric: If the input tensor is quantized symmetrically.
use_fp: Use floating point.
add_scale_eps: Add eps value or replace zero value by 1 to avoid division by
zero.
per_channel: use per-channel clipping optimization.
random_rounding: round with uniform random.
key: rng key for rounding.
Returns:
A tuple of quantized tensor, quantization scale
and quantization zero point (optional).
"""
min_value, max_value = get_min_max(bits, use_fp=use_fp)
if use_symmetric:
bound = jnp.max(jnp.abs(t), axis=contract_dims, keepdims=True)
scale_bound = max_value
else:
t_max = jnp.max(t, axis=contract_dims, keepdims=True)
t_min = jnp.min(t, axis=contract_dims, keepdims=True)
bound = t_max - t_min
scale_bound = max_value - min_value
if percentile < 1.0:
bound = jnp.multiply(bound, percentile)
elif optimization_on_bound:
bound = get_best_bound(
t, bound, min_value, max_value, p_value, per_channel=per_channel
)
scale = bound / scale_bound
if add_scale_eps:
# Add epsilon to avoid divide-by-zero.
scale = scale + jnp.finfo(t.dtype).eps
else:
scale = jnp.where(scale == 0.0, 1.0, scale)
if use_symmetric:
zp = None
t = jnp.divide(t, scale)
else:
zp = min_value - t_min / scale
t = jnp.divide(t, scale) + zp
zp = jnp.multiply(scale, zp)
if use_fp:
# No need to round.
t = jnp.clip(t, min_value, max_value).astype(jnp.float8_e4m3fn)
# TODO: refactor to remove this logic.
t = jax.lax.bitcast_convert_type(t, new_dtype=jnp.int8)
else:
if need_gradient:
t = pass_through(t, jnp.round)
t = jnp.clip(t, min_value, max_value)
else:
if random_rounding:
t = t + jax.random.uniform(
key=key, shape=t.shape, minval=-0.5, maxval=0.5
)
t = jnp.round(t)
container_dtype = (
jnp.int8 if bits <= 8 else jnp.int16 if bits <= 16 else jnp.int32
)
t = jnp.clip(t, min_value, max_value).astype(container_dtype)
return t, scale, zp
def quantize_tensor(
var: np.ndarray,
axis: List[int],
factor: float = 1.0,
sym: bool = True,
number_bits: int = 8,
use_fp: bool = False,
add_scale_eps: bool = False,
optimization_on_bound: bool = False,
p_value: float = 1.0,
per_channel: bool = False,
block_size: int = 0,
use_mse_quant: bool = False,
) -> Union[
Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]
]:
"""Quantize a tensor.
Args:
var: The variable to be quantized.
axis: The axis along which variable will be quantized.
factor: The clipping factor.
sym: Symmetric or asymmetric quantize the variable.
number_bits: Number of bits for quantized value.
use_fp: do fp with number of bits (i.e. fp8)
add_scale_eps: add epsilon to scale to avoid division by zero, else it will
replace zero scale by 1.
optimization_on_bound: If p-mean bound optimizer is used.
p_value: Exponent of the p-mean error metric. Default to 1.0 which is MAE.
per_channel: use per-channel clipping optimization.
block_size: block size for sub-channel quantization. Defaults to 0, which
means off.
use_mse_quant: Whether to use MSE quantization.
Returns:
Quantized tensors, along with scales and zero point.
"""
# TODO: support jnp.float8_e5m2
assert number_bits == 8 or number_bits == 4
jnp_var = jnp.asarray(var)
# When using sub-channel, the contracting dim is split into a sub-channel
# dim followed by the block dim. Therefore the contracting dim
# (quantize_axis) should increment by one, and the corresponding pack_dim
# should also increment by one.
if block_size > 0:
# We hard-code in a special case for [0, 1]. It has not been thoroughly
# tested or verified yet, so we log a warning.
if len(axis) == 2 and axis == [0, 1]:
print('Using experimental blockwise int4 axis=[0, 1] compression')
jnp_var = jnp.reshape(jnp_var, [-1, jnp_var.shape[-1]])
axis = [0]
shape = list(jnp_var.shape)
assert len(axis) == 1, 'Only support 1D sub-channel quantization'
sub_channels, rem = divmod(shape[axis[0]], block_size)
assert rem == 0
shape.insert(axis[0], sub_channels)
axis[0] += 1
shape[axis[0]] = block_size
jnp_var = jnp.reshape(jnp_var, shape)
if use_mse_quant:
qvar, scale, zp = mse_reduce_precision(
jnp_var,
contract_dims=axis,
need_gradient=False,
bits=number_bits,
use_symmetric=sym,
use_fp=False,
)
else:
qvar, scale, zp = reduce_precision(
jnp_var,
contract_dims=axis,
need_gradient=False,
bits=number_bits,
optimization_on_bound=optimization_on_bound,
percentile=factor,
use_symmetric=sym,
use_fp=use_fp,
add_scale_eps=add_scale_eps,
p_value=p_value,
per_channel=per_channel,
)
if sym:
return np.array(qvar), np.array(jnp.squeeze(scale, axis=axis)) # pytype: disable=wrong-arg-types # jnp-type
else:
return (
np.array(qvar),
# CAVEAT: the following squeezes should squeeze along the quantization
# axis only.
np.array(jnp.squeeze(scale)),
np.array(jnp.squeeze(zp)),
)
def pack_4bit(
x: np.ndarray, pack_dim: int, packed_dtype: jnp.dtype = jnp.int32
) -> np.ndarray:
"""Pack int8 or uint8 tensor where its values are actually int4 or uint4, to int32 or int8 nibble format along pack_dim.
Args:
x: Original int8 or uint8 tensor to pack.
pack_dim: Dimension to pack along. x.shape[pack_dim] must be divisible by 8,
when packed_dtype is int32 and divisible by 2 when target_type is int8.
Also pack_dim must be < x.ndim - 1.
packed_dtype: Target type to pack to, int32 or int8.
Returns:
int32 or int8 packed tensor where the pack_dim size is dividend by 8
from the original tensor x.
"""
x = jnp.asarray(x)
if packed_dtype == jnp.int8 and x.dtype == jnp.uint8:
# It doesn't make sense to pack uint8 numbers into int4 as we'll
# the range overlap between uint8 and int4 is [0..7].
raise ValueError(
'only int8 input dtype is supported when packing into int8. '
f'Given {x.dtype}'
)
if x.dtype != jnp.int8 and x.dtype != jnp.uint8:
raise ValueError(
f'input dtype must be either int8 or uint8. Given {x.dtype}'
)
if pack_dim >= x.ndim - 1:
raise ValueError(
f'pack_dim must be < input ndim - 1. input shape {x.shape} and pack_dim'
f' {pack_dim}'
)
if packed_dtype != jnp.int32 and packed_dtype != jnp.int8:
raise ValueError(
f'packed_dtype must be either int32 or int8. Given {packed_dtype}'
)
if packed_dtype == jnp.int32 and x.shape[pack_dim] % 8 != 0:
raise ValueError(
'input shape[pack_dim] must be divisible by 8 when target_type '
f'is int32. Given shape {x.shape}'
)
if packed_dtype == jnp.int8 and x.shape[pack_dim] % 2 != 0:
raise ValueError(
'input shape[pack_dim] must be divisible by 2 when target_type '
f'is int8. Given shape {x.shape}'
)
int4s_per_packed_type = 8 if packed_dtype == jnp.int32 else 2
rep_shape = list(x.shape)
rep_shape.insert(pack_dim + 1, int4s_per_packed_type)
rep_shape[pack_dim] //= int4s_per_packed_type
shifts = lax.broadcasted_iota(packed_dtype, rep_shape, pack_dim + 1)
shifts <<= 2
# Promote x to packed_dtype
x = x & jnp.array(0x0F, packed_dtype)
x = lax.reshape(x, rep_shape)
x = x << shifts
x = lax.reduce(x, jnp.array(0x0, packed_dtype), lax.add, [pack_dim + 1])
return np.asarray(x)
def pack_2bit_into_int8(x: np.ndarray, pack_dim: int) -> np.ndarray:
"""Pack (u)int8 tensor with int2 or uint2 values into int8 along pack_dim.
Args:
x: Original int8 or uint8 tensor to pack.
pack_dim: Dimension to pack along. x.shape[pack_dim] must be divisible by 4.
Also pack_dim must be < x.ndim - 1.
Returns:
int8 packed tensor where the pack_dim size is divided by 4 from the original
tensor x.
"""
x = jnp.asarray(x)
packed_dtype = jnp.int8
bits = 2
bits_per_packed_type = 8 // bits # 4 int2 values per int8
if x.dtype != jnp.int8 and x.dtype != jnp.uint8:
raise ValueError(
f'input dtype must be either int8 or uint8. Given {x.dtype}'
)
if pack_dim >= x.ndim - 1:
raise ValueError(
f'pack_dim must be < input ndim - 1. input shape {x.shape} and pack_dim'
f' {pack_dim}'
)
if x.shape[pack_dim] % bits_per_packed_type != 0:
raise ValueError(
f'input shape[pack_dim] must be divisible by {bits_per_packed_type}'
f' when packing {bits}-bit values into {packed_dtype}. Given shape'
f' {x.shape}'
)
rep_shape = list(x.shape)
rep_shape.insert(pack_dim + 1, bits_per_packed_type)
rep_shape[pack_dim] //= bits_per_packed_type
# Create shifts: [0, 1, 2, 3]
shifts = lax.broadcasted_iota(packed_dtype, rep_shape, pack_dim + 1)
# Adjust shifts for 2-bit packing: [0, 2, 4, 6]
shifts <<= bits - 1 # Equivalent to shifts * 2
# Promote x to packed_dtype and mask to keep only the lowest 2 bits (0x03)
x = x & jnp.array(0x03, packed_dtype)
x = lax.reshape(x, rep_shape)
x = x << shifts
x = lax.reduce(x, jnp.array(0x0, packed_dtype), lax.add, [pack_dim + 1])
return np.asarray(x)
def update_to_uint4(
qx: np.ndarray, scale: np.ndarray, zp: Optional[np.ndarray] = None
):
"""Updates the quantized weights from int4 to uint4.
This is a conversion function designed for XNNPack as it expects the 4-bit
quantized weight to be represented differently from the original Pax setting.
Specifically, the differences are:
1) The dynamic range of weight values: int4 (Pax) vs. uint4 (XNNPack).
2) The dynamic range of zero-point: float (Pax) vs. uint4 (XNNPack).
3) The number of zero-point: per-channel (Pax) vs. per-tensor (XNNPack).
Args:
qx: np.array of shape [..., channel], which is the quantized weight values
from Pax in the shape of. The values are in the dynamic range of int4 but
are hosted as int8 type. Note that if the first dimension is 3, it means
the qkv matrices are concatenated together and should be treated
differently.
scale: np.array of shape [1(3), channel] as np.float type, which are the
scaling factors for dequantization per channel.
zp: (optional) np.array of shape [1 (or 3), channel] as np.float type, which
are the zero points for dequantization per channel.
Returns:
A tuple (qx, scale, zp):
qx: The updated np.array of shape [..., channel] as np.int8 type with
updated dynamic range as uint4 (with 8 as the default zero points).
scale: Same as the input scale.
zp: (optional) np.array of shape [1 (or 3)] as np.int8 type with the
updated zero point values in the dynamic range as uint4.
"""
if qx.dtype != np.int8 or ('float' not in str(scale.dtype)):
raise ValueError(
'Unexpected dtype qx:' + str(qx.dtype) + ' scale:' + str(scale.dtype)
)
scale = scale.astype(np.float32)
def get_new_zp(old_zp):
new_zp = old_zp / (scale + np.finfo(np.float32).eps)
per_tensor_zp = np.mean(new_zp)
per_tensor_zp = per_tensor_zp.astype(np.int8) + _UINT4_ZP
return per_tensor_zp
if zp is not None:
if qx.shape[0] == 3:
per_tensor_zp = np.stack([get_new_zp(szp) for szp in zp], axis=0)
else:
per_tensor_zp = get_new_zp(zp)
else:
per_tensor_zp = (
_UINT4_ZP * np.ones(shape=(3)) if qx.shape[0] == 3 else _UINT4_ZP
)
qx = qx + _UINT4_ZP
return qx, scale, np.array(per_tensor_zp, dtype=np.int32)
def max_length_axis(shape, axis):
"""Get the axis with the maximum length."""
if axis is None:
axis = range(len(shape))
elif isinstance(axis, int):
axis = [axis]
axis = tuple(a % len(shape) for a in axis)
axis = sorted(axis, reverse=True) # Prefers larger dim if there is a tie.
max_len_ind = max(enumerate([shape[a] for a in axis]), key=lambda x: x[1])[0]
return axis[max_len_ind]
def dequantize_blockwise(values, scales, block_size):
"""Dequantizes a given blockwise-compressed tensor."""
contracted_dims = []
if len(values.shape) != len(scales.shape):
raise ValueError('Values and scales are not the same size')
# Find the contracted dimensions.
for i in range(len(values.shape)):
if values.shape[i] != scales.shape[i]:
contracted_dims.append(i)
# Make sure decompressing results in an exact shape match.
test_shape = list(scales.shape)
for i in contracted_dims:
test_shape[i] *= block_size
if tuple(test_shape) != values.shape:
raise ValueError('Expanded dimensions did not match')
# We need to repeat the scale tensor along the target dimension to match the
# shape of the quantized value tensor.
target_dim = max_length_axis(values.shape, contracted_dims)
scales = jnp.repeat(scales, block_size, axis=target_dim)
return values.astype(jnp.int32).astype(jnp.float32) * scales
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/quantization_util_test.py
+292
View File
@@ -0,0 +1,292 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Tests for quantization_util."""
from absl.testing import absltest
import jax
from jax import numpy as jnp
import numpy as np
from mediapipe.tasks.python.genai.converter import quantization_util
_dtype = lambda x: getattr(x, 'dtype', None) or np.asarray(x).dtype
class TestCase(absltest.TestCase):
def assertAllClose(
self, x, y, check_dtypes=True, rtol=1e-5, atol=1e-5, **kwargs
):
"""Wrapper for np.testing.assert_allclose()."""
x = np.asarray(x)
y = np.asarray(y)
if check_dtypes:
self.assertDtypesMatch(x, y)
x = x.astype(np.float32) if x.dtype == jnp.bfloat16 else x
y = y.astype(np.float32) if y.dtype == jnp.bfloat16 else y
np.testing.assert_allclose(x, y, rtol=rtol, atol=atol, **kwargs)
def assertDtypesMatch(self, x, y):
self.assertEqual(
jax.dtypes.canonicalize_dtype(_dtype(x)),
jax.dtypes.canonicalize_dtype(_dtype(y)),
)
class Quantize8BTest(TestCase):
def test_quantize_symmetric(self):
inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
qx, scale = quantization_util.quantize_tensor(inputs, axis=[1])
self.assertAllClose(
qx, np.array([[28, 72, 127, 67], [6, -48, 105, 127]], dtype=np.int8)
)
self.assertAllClose(
scale, np.array([0.04330709, 0.03149606], dtype=np.float32)
)
def test_quantize_symmetric_with_dimension_size_one_unquantized(self):
# inputs shape: (2, 1, 4), quantization axis 2.
inputs = np.array([[[1.2, 3.1, 5.5, 2.9]], [[0.2, -1.5, 3.3, 4.0]]])
qx, scale = quantization_util.quantize_tensor(inputs, axis=[2])
self.assertAllClose(
qx, np.array([[[28, 72, 127, 67]], [[6, -48, 105, 127]]], dtype=np.int8)
)
# expected scale shape: (2, 1)
self.assertAllClose(
scale, np.array([[0.04330709], [0.03149606]], dtype=np.float32)
)
def test_quantize_asymmetric(self):
inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
qx, scale, zp = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False
)
self.assertAllClose(
qx,
np.array([[-128, -15, 127, -27], [-49, -128, 95, 127]], dtype=np.int8),
)
self.assertAllClose(scale, np.array([0.016863, 0.021569], dtype=np.float32))
self.assertAllClose(zp, np.array([-3.358431, -1.260784], dtype=np.float32))
class Quantize8BFPTest(TestCase):
def test_quantize_symmetric(self):
inputs = np.array([[1.0, 2.0, 5.5, 2.9], [0.02, -0.01, 3.3, 4.0]])
qx, scale = quantization_util.quantize_tensor(inputs, axis=[1], use_fp=True)
self.assertAllClose(
qx,
np.array([[106, 114, 126, 119], [65, -71, 124, 126]], dtype=np.int8),
)
self.assertAllClose(
scale, np.array([0.01227679, 0.00892857], dtype=np.float32)
)
def test_quantize_symmetric_with_dimension_size_one_unquantized(self):
# inputs shape: (2, 1, 4), quantization axis 2.
inputs = np.array([[[1.0, 2.0, 5.5, 2.9]], [[0.02, -0.01, 3.3, 4.0]]])
qx, scale = quantization_util.quantize_tensor(inputs, axis=[2], use_fp=True)
self.assertAllClose(
qx,
np.array(
[[[106, 114, 126, 119]], [[65, -71, 124, 126]]], dtype=np.int8
),
)
# expected scale shape: (2, 1)
self.assertAllClose(
scale, np.array([[0.01227679], [0.00892857]], dtype=np.float32)
)
def test_quantize_asymmetric(self):
inputs = np.array([[-1.0, -2.0, -2.01, 2.9], [0.02, -0.15, 3.3, 4.0]])
qx, scale, zp = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False, use_fp=True
)
self.assertAllClose(
qx,
np.array([[-8, -2, -2, 126], [-3, -2, 121, 126]], dtype=np.int8),
)
self.assertAllClose(
scale, np.array([0.00547991, 0.0046317], dtype=np.float32)
)
self.assertAllClose(
zp, np.array([-0.4449999, -1.9250002], dtype=np.float32)
)
def test_quantize_add_scale_eps(self):
inputs = np.array([[0.0, 0.0, 0.0, 0.0], [-4.0, -4.0, -4.0, -4.0]])
_, scale, _ = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False, use_fp=True, add_scale_eps=True
)
self.assertAllClose(
scale, np.array([np.finfo(np.float32).eps, np.finfo(np.float32).eps])
)
_, scale, _ = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False, use_fp=True, add_scale_eps=False
)
self.assertAllClose(scale, np.array([1.0, 1.0]))
class Quantize4BTest(TestCase):
def test_quantize_symmetric(self):
inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
qx, scale = quantization_util.quantize_tensor(
inputs, axis=[1], number_bits=4
)
self.assertAllClose(
qx, np.array([[2, 4, 7, 4], [0, -3, 6, 7]], dtype=np.int8)
)
self.assertAllClose(
scale, np.array([0.78571427, 0.5714286], dtype=np.float32)
)
def test_quantize_symmetric_with_dimension_size_one_unquantized(self):
# inputs shape: (2, 1, 4), quantization axis 2.
inputs = np.array([[[1.2, 3.1, 5.5, 2.9]], [[0.2, -1.5, 3.3, 4.0]]])
qx, scale = quantization_util.quantize_tensor(
inputs, axis=[2], number_bits=4
)
self.assertAllClose(
qx, np.array([[[2, 4, 7, 4]], [[0, -3, 6, 7]]], dtype=np.int8)
)
# expected scale shape: (2, 1)
self.assertAllClose(
scale, np.array([[0.78571427], [0.5714286]], dtype=np.float32)
)
def test_quantize_asymmetric(self):
inputs = np.array([[1.2, 3.1, 5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
qx, scale, zp = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False, number_bits=4
)
self.assertAllClose(
qx,
np.array([[-8, -1, 7, -2], [-3, -8, 5, 7]], dtype=np.int8),
)
self.assertAllClose(
scale, np.array([0.2866667, 0.36666667], dtype=np.float32)
)
self.assertAllClose(
zp, np.array([-3.4933336, -1.4333334], dtype=np.float32)
)
class QuantizationUtilTest(TestCase):
def test_update_to_uint4_sym(self):
inputs = np.array([[1.2, 3.1, -5.5, 2.9], [0.2, -1.5, 3.3, 4.0]])
qx, scale = quantization_util.quantize_tensor(
inputs, axis=[1], sym=True, number_bits=4
)
dequant_from_int4 = qx * np.expand_dims(scale, -1)
qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale)
self.assertEmpty(zp_n.shape) # A scalar numpy array.
dequant_from_uint4 = np.expand_dims(scale_n, -1) * (qx_n - zp_n)
np.testing.assert_allclose(dequant_from_int4, dequant_from_uint4)
def test_update_to_uint4_sym_combined(self):
inputs = np.array(
[[-1.2, 3.5, -6.2, 1.7], [1.2, 3.1, -5.5, 2.9], [0.2, -1.5, 3.3, 4.0]]
)
qx, scale = quantization_util.quantize_tensor(
inputs, axis=[1], sym=True, number_bits=4
)
dequant_from_int4 = qx * np.expand_dims(scale, -1)
qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale)
self.assertEqual(zp_n.shape[0], 3)
dequant_from_uint4 = np.expand_dims(scale_n, -1) * (
qx_n - np.expand_dims(zp_n, -1)
)
np.testing.assert_allclose(dequant_from_int4, dequant_from_uint4)
def test_update_to_uint4_asym(self):
inputs = np.array([[1.0, 8.0, -3.0, 2.0], [-3.0, 2.0, 1.0, 8.0]])
qx, scale, zp = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False, number_bits=4
)
qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale, zp)
expected_dequant = np.array([
[0.0, 7.333333, -3.666667, 1.466667],
[-3.666667, 1.466667, 0.0, 7.333333],
])
dequant_from_uint4 = np.expand_dims(scale_n, -1) * (qx_n - zp_n)
np.testing.assert_allclose(dequant_from_uint4, expected_dequant, rtol=1e-05)
def test_update_to_uint4_asym_combined(self):
inputs = np.array(
[[1.0, 8.0, -3.0, 2.0], [-3.0, 2.0, 1.0, 8.0], [2.0, 1.0, 8.0, -3.0]]
)
qx, scale, zp = quantization_util.quantize_tensor(
inputs, axis=[1], sym=False, number_bits=4
)
qx_n, scale_n, zp_n = quantization_util.update_to_uint4(qx, scale, zp)
self.assertEqual(zp_n.shape[0], 3)
expected_dequant = np.array([
[0.0, 7.333333, -3.666667, 1.466667],
[-3.666667, 1.466667, 0.0, 7.333333],
[1.466667, 0.0, 7.333333, -3.666667],
])
dequant_from_uint4 = np.expand_dims(scale_n, -1) * (
qx_n - np.expand_dims(zp_n, -1)
)
np.testing.assert_allclose(dequant_from_uint4, expected_dequant, rtol=1e-05)
class QuantizeMseTest(TestCase):
def test_mse_reduce_precision_4bit(self):
inputs = np.array([[-1.0, 1.0, 0.0], [1.0, 2.0, 3.0]], dtype=np.float32)
qx, scale = quantization_util.quantize_tensor(
inputs, axis=[1], number_bits=4, use_mse_quant=True
)
# For 4 bits, maxq = 15, zero_mse_quant = (15 - 1) / 2 = 7.
# The scale multiplier is 0.37755.
# For inputs row 1 ([-1.0, 1.0, 0.0]):
# rms = sqrt(mean((-1)^2 + 1^2 + 0^2)) = sqrt(2/3) approx 0.8164966
# scale = 0.37755 * 0.8164966 = 0.3079979
# t_shifted(-1.0) = -1.0 / 0.3079979 + 7 = -3.2461 + 7 = 3.7539
# -> round 4 -> qvar = -3
# t_shifted(1.0) = 1.0 / 0.3079979 + 7 = 3.2461 + 7 = 10.2461
# -> round 10 -> qvar = 3
# t_shifted(0.0) = 0.0 / 0.3079979 + 7 = 7.0 -> round 7 -> qvar = 0
# For inputs row 2 ([1.0, 2.0, 3.0]):
# rms = sqrt(mean(1^2 + 2^2 + 3^2)) = sqrt(14/3) approx 2.1602469
# scale = 0.37755 * 2.1602469 = 0.8159988
# t_shifted(1.0) = 1.0 / 0.8159988 + 7 = 1.2255 + 7 = 8.2255
# -> round 8 -> qvar = 1
# t_shifted(2.0) = 2.0 / 0.8159988 + 7 = 2.4509 + 7 = 9.4509
# -> round 9 -> qvar = 2
# t_shifted(3.0) = 3.0 / 0.8159988 + 7 = 3.6764 + 7 = 10.6764
# -> round 11 -> qvar = 4
self.assertAllClose(qx, np.array([[-3, 3, 0], [1, 2, 4]], dtype=np.int8))
self.assertAllClose(scale, np.array([0.308268, 0.815601], dtype=np.float32))
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/safetensors_converter.py
+635
View File
@@ -0,0 +1,635 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""CkptLoader implementation for loading the Safetensors."""
import array
from typing import Iterator
import enum
import glob
import json
import os
from typing import List, Optional
import numpy as np
from mediapipe.tasks.python.genai.converter import converter_base
from mediapipe.tasks.python.genai.converter import external_dependencies
torch = external_dependencies.torch
DTYPE_MAP = {
"F16": torch.float16,
"BF16": torch.bfloat16,
"F32": torch.float32,
}
class _SafetensorsShardReader:
"""Reads a single safetensors shard."""
_HEAD_BYTES = 8
def __init__(self, shard_path: str):
self._shard_path = shard_path
if not os.path.exists(self._shard_path):
raise ValueError(f"{self._shard_path} does not exists.")
with open(self._shard_path, "rb") as f:
head_bytes = f.read(self._HEAD_BYTES)
metadata_bytes_num = np.frombuffer(head_bytes, dtype=np.uint64)[0]
metadata_bytes = f.read(metadata_bytes_num)
self.layers_info = json.loads(metadata_bytes)
self.metadata_bytes_num = metadata_bytes_num
def read_tensor_as_numpy(self, tensor_name) -> np.ndarray:
"""Reads a tensor from the model file as a numpy array with np.float32 type."""
tensor_info = self.layers_info[tensor_name]
with open(self._shard_path, "rb") as f:
shape = tensor_info["shape"]
dtype = tensor_info["dtype"]
if dtype not in DTYPE_MAP:
raise ValueError(f"{dtype} is not supported.")
data_offsets = tensor_info["data_offsets"]
f.seek(int(self._HEAD_BYTES + self.metadata_bytes_num + data_offsets[0]))
tensor_bytes = f.read(data_offsets[1] - data_offsets[0])
raw_tensor = torch.frombuffer(
array.array("b", tensor_bytes), dtype=DTYPE_MAP[dtype]
).reshape(shape)
return raw_tensor.float().t().contiguous().numpy()
def get_tensor_names(self) -> List[str]:
names = list(self.layers_info.keys())
if "__metadata__" in names:
names.remove("__metadata__")
return names
class _SafetensorsReader:
"""Reads all the safetensors shards."""
def __init__(self, ckpt_path: str):
shards = []
if os.path.isdir(ckpt_path):
# Read all safetensors files within checkpoint
for shard_path in glob.glob(os.path.join(ckpt_path, "*.safetensors")):
shards.append(_SafetensorsShardReader(shard_path))
else:
# Assume the ckpt_path is a file or a file pattern to match.
for shard_path in glob.glob(ckpt_path):
shards.append(_SafetensorsShardReader(shard_path))
assert shards is not None
self._ckpt_path = ckpt_path
self._tensors_map = {}
for shard in shards:
tensor_names = shard.get_tensor_names()
for tensor_name in tensor_names:
if tensor_name in self._tensors_map:
raise ValueError(f"Duplicate tensor name: {tensor_name}")
self._tensors_map[tensor_name] = shard
def get_tensor_names(self) -> List[str]:
return list(self._tensors_map.keys())
def read_tensor_as_numpy(self, tensor_name: str) -> np.ndarray:
return self._tensors_map[tensor_name].read_tensor_as_numpy(tensor_name)
class LayerType(enum.Enum):
"""Enum for layer type."""
NONE = 0
ATTENTION = 1 # Layer is part of the attention module.
FEEDFORWARD = 2 # Layer is part of the feedforward module in the Transformer.
EMBEDDING = 3 # Layer is the embedding lookup or final projection layer.
LAYER_NORM = (
4 # Layer is layer normalization before and after attention layer.
)
LORA = 5 # Layer is LoRA weights augmented on the base model layers.
SKIP = 6 # Log a warning, and remove this not-yet-supported layer entirely.
@classmethod
def get_layer_type(cls, layer_name: str):
"""Gets the layer type of the given layer name."""
ffn_layers = [
"mlp",
]
attn_layers = [
"self_attn",
]
emb_layers = [
"embed_tokens",
"embedder",
"lm_head",
]
layer_norms = [
"input_layernorm",
"post_attention_layernorm",
"final_layernorm",
"model.norm.weight",
"pre_feedforward_layernorm",
"post_feedforward_layernorm",
]
lora_layers = ["lora"]
skip_layers = [
"multi_modal_projector",
"vision_tower",
]
if any(sub_name in layer_name for sub_name in skip_layers):
return LayerType.SKIP
if any(sub_name in layer_name for sub_name in lora_layers):
return LayerType.LORA
if any(sub_name in layer_name for sub_name in attn_layers):
return LayerType.ATTENTION
if any(sub_name in layer_name for sub_name in ffn_layers):
return LayerType.FEEDFORWARD
if any(sub_name in layer_name for sub_name in emb_layers):
return LayerType.EMBEDDING
if any(sub_name in layer_name for sub_name in layer_norms):
return LayerType.LAYER_NORM
else:
return LayerType.NONE
class StablelmMapper(converter_base.LayerActionMapperBase):
"""LayerActionMapper for handling the StableLM model."""
def __init__(
self,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
backend: str,
reader: _SafetensorsReader,
):
super().__init__(
is_symmetric=is_symmetric,
attention_quant_bits=attention_quant_bits,
feedforward_quant_bits=feedforward_quant_bits,
embedding_quant_bits=embedding_quant_bits,
backend=backend,
)
self._reader = reader
def map_to_actions(
self, layer_name: str
) -> Optional[List[converter_base.QuantizationAction]]:
"""Map the given layer name to actions."""
tensor_value = self._reader.read_tensor_as_numpy(layer_name)
quantize_axis = None
quantize_bits = None
layer_type = LayerType.get_layer_type(layer_name)
if layer_type != LayerType.LAYER_NORM and layer_name.endswith(".weight"):
quantize_axis = [0]
if layer_type == LayerType.FEEDFORWARD:
quantize_bits = self._feedforward_quant_bits
elif layer_type == LayerType.ATTENTION:
quantize_bits = self._attention_quant_bits
if self._backend == "cpu" and ".o_proj." in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
elif layer_type == LayerType.EMBEDDING:
quantize_bits = self._embedding_quant_bits
if self._backend == "cpu" and ".embed_tokens." in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
target_name = self.update_target_name(layer_name)
actions = [
converter_base.QuantizationAction(
tensor_name=layer_name,
tensor_value=tensor_value,
target_name=target_name,
quantize_axis=quantize_axis,
quantize_bits=quantize_bits,
pack_dim=0,
)
]
return actions
def update_target_name(self, target_name: str) -> str:
"""Updates the target name to match the tensor name convention."""
target_name = target_name.replace(
"model.layers.", "params.lm.transformer.x_layers_"
)
target_name = target_name.replace("mlp.up_proj", "ff_layer.ffn_layer1")
target_name = target_name.replace("mlp.down_proj", "ff_layer.ffn_layer2")
target_name = target_name.replace(
"mlp.gate_proj", "ff_layer.ffn_layer1_gate"
)
target_name = target_name.replace("input_layernorm", "pre_layer_norm")
target_name = target_name.replace(
"pre_layer_norm.weight", "pre_layer_norm.scale"
)
if self._backend == "cpu":
target_name = target_name.replace(
"post_attention_layernorm", "ff_layer.pre_layer_norm"
)
target_name = target_name.replace(
"ff_layer.pre_layer_norm.weight", "ff_layer.pre_layer_norm.scale"
)
else:
target_name = target_name.replace(
"post_attention_layernorm", "post_layer_norm"
)
target_name = target_name.replace(
"post_layer_norm.weight", "post_layer_norm.scale"
)
target_name = target_name.replace("self_attn.q_proj", "self_attention.q")
target_name = target_name.replace("self_attn.k_proj", "self_attention.k")
target_name = target_name.replace("self_attn.v_proj", "self_attention.v")
target_name = target_name.replace("self_attn.o_proj", "self_attention.post")
target_name = target_name.replace(
"model.embed_tokens", "params.lm.token_embedding"
)
target_name = target_name.replace("model.norm", "params.lm.final_ln")
target_name = target_name.replace("final_ln.weight", "final_ln.scale")
target_name = target_name.replace("lm_head", "params.lm.softmax.logits_ffn")
target_name = target_name.replace(".weight", ".w")
return target_name
class PhiMapper(converter_base.LayerActionMapperBase):
"""LayerActionMapper for handling the Phi model."""
def __init__(
self,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
backend: str,
reader: _SafetensorsReader,
):
super().__init__(
is_symmetric=is_symmetric,
attention_quant_bits=attention_quant_bits,
feedforward_quant_bits=feedforward_quant_bits,
embedding_quant_bits=embedding_quant_bits,
backend=backend,
)
self._reader = reader
def map_to_actions(
self, layer_name: str
) -> Optional[List[converter_base.QuantizationAction]]:
"""Map the given layer name to actions."""
tensor_value = self._reader.read_tensor_as_numpy(layer_name)
quantize_axis = None
quantize_bits = None
layer_type = LayerType.get_layer_type(layer_name)
if (
layer_type != LayerType.LAYER_NORM
and layer_name.endswith(".weight")
and layer_type != LayerType.LORA
):
quantize_axis = [0]
if layer_type == LayerType.FEEDFORWARD:
quantize_bits = self._feedforward_quant_bits
elif layer_type == LayerType.ATTENTION:
quantize_bits = self._attention_quant_bits
if self._backend == "cpu" and ".dense." in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
elif layer_type == LayerType.EMBEDDING:
quantize_bits = self._embedding_quant_bits
if self._backend == "cpu" and ".embed_tokens." in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
target_name = self.update_target_name(layer_name)
actions = [
converter_base.QuantizationAction(
tensor_name=layer_name,
tensor_value=tensor_value,
target_name=target_name,
quantize_axis=quantize_axis,
quantize_bits=quantize_bits,
pack_dim=0,
)
]
return actions
def update_target_name(self, target_name: str) -> str:
"""Updates the target name to match the tensor name convention."""
target_name = target_name.replace("base_model.model.", "")
target_name = target_name.replace(
"model.layers.", "params.lm.transformer.x_layers_"
)
layer_type = LayerType.get_layer_type(target_name)
if layer_type == LayerType.FEEDFORWARD:
target_name = target_name.replace(".weight", ".linear.w")
target_name = target_name.replace(".bias", ".bias.b")
target_name = target_name.replace("mlp.fc1", "ff_layer.ffn_layer1")
target_name = target_name.replace("mlp.fc2", "ff_layer.ffn_layer2")
elif layer_type == LayerType.ATTENTION:
target_name = target_name.replace(".weight", ".linear.w")
target_name = target_name.replace(".bias", ".bias.b")
target_name = target_name.replace("self_attn.q_proj", "self_attention.q")
target_name = target_name.replace("self_attn.k_proj", "self_attention.k")
target_name = target_name.replace("self_attn.v_proj", "self_attention.v")
target_name = target_name.replace(
"self_attn.dense", "self_attention.post"
)
elif layer_type == LayerType.EMBEDDING:
target_name = target_name.replace(
"model.embed_tokens", "params.lm.token_embedding"
)
target_name = target_name.replace(
"lm_head", "params.lm.softmax.logits_ffn"
)
target_name = target_name.replace(
"logits_ffn.weight", "logits_ffn.linear.w"
)
target_name = target_name.replace("logits_ffn.bias", "logits_ffn.bias.b")
elif layer_type == LayerType.LAYER_NORM:
target_name = target_name.replace("input_layernorm", "pre_layer_norm")
target_name = target_name.replace(
"pre_layer_norm.weight", "pre_layer_norm.scale"
)
target_name = target_name.replace(
"model.final_layernorm", "params.lm.final_ln"
)
target_name = target_name.replace("final_ln.weight", "final_ln.scale")
target_name = target_name.replace(".weight", ".w")
# For LoRA weights
if "post" in target_name:
target_name = target_name.replace("lora_A.linear.w", "w_prime_right")
target_name = target_name.replace("lora_B.linear.w", "w_prime_left")
else:
target_name = target_name.replace("lora_A.linear.w", "w_prime_left")
target_name = target_name.replace("lora_B.linear.w", "w_prime_right")
return target_name
class GemmaMapper(converter_base.LayerActionMapperBase):
"""LayerActionMapper for handling the Gemma model."""
def __init__(
self,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
backend: str,
reader: _SafetensorsReader,
is_v2: bool,
):
super().__init__(
is_symmetric=is_symmetric,
attention_quant_bits=attention_quant_bits,
feedforward_quant_bits=feedforward_quant_bits,
embedding_quant_bits=embedding_quant_bits,
backend=backend,
)
self._reader = reader
self._is_v2 = is_v2
def map_to_actions(
self, layer_name: str
) -> Optional[List[converter_base.QuantizationAction]]:
"""Map the given layer name to actions."""
layer_type = LayerType.get_layer_type(layer_name)
if layer_type == LayerType.SKIP:
print("Unsupported Gemma layer: %s, skipping." % layer_name)
return
tensor_value = self._reader.read_tensor_as_numpy(layer_name)
quantize_axis = None
quantize_bits = None
if (
layer_type != LayerType.LAYER_NORM
and layer_name.endswith(".weight")
and layer_type != LayerType.LORA
):
quantize_axis = [0]
if layer_type == LayerType.FEEDFORWARD:
quantize_bits = self._feedforward_quant_bits
elif layer_type == LayerType.ATTENTION:
quantize_bits = self._attention_quant_bits
if "o_proj" in layer_name:
tensor_value = np.transpose(tensor_value)
quantize_axis = [1]
if "q_norm" in layer_name:
# TODO: Look into quantization for q_norm and k_norm
quantize_axis = None
quantize_bits = None
if "k_norm" in layer_name:
quantize_axis = None
quantize_bits = None
elif layer_type == LayerType.EMBEDDING:
quantize_bits = self._embedding_quant_bits
target_name = self.update_target_name(layer_name)
actions = [
converter_base.QuantizationAction(
tensor_name=layer_name,
tensor_value=tensor_value,
target_name=target_name,
quantize_axis=quantize_axis,
quantize_bits=quantize_bits,
pack_dim=0,
)
]
return actions
def update_target_name(self, target_name: str) -> str:
"""Updates the target name to match the tensor name convention."""
# For removing multimodality stack from Gemma3-4B
target_name = target_name.replace("language_model.", "")
target_name = target_name.replace("base_model.model.", "")
target_name = target_name.replace(
"model.layers.", "params.lm.transformer.x_layers_"
)
target_name = target_name.replace("mlp.up_proj", "ff_layer.ffn_layer1")
target_name = target_name.replace("mlp.down_proj", "ff_layer.ffn_layer2")
target_name = target_name.replace(
"mlp.gate_proj", "ff_layer.ffn_layer1_gate"
)
target_name = target_name.replace("input_layernorm", "pre_layer_norm")
target_name = target_name.replace(
"pre_layer_norm.weight", "pre_layer_norm.scale"
)
# Gemma and Gemma2/Gemma3 differ slightly in their use of the
# "post_attention_layernorm" tensor name.
if self._is_v2:
target_name = target_name.replace(
"post_attention_layernorm", "post_layer_norm"
)
else:
target_name = target_name.replace(
"post_attention_layernorm", "ff_layer.pre_layer_norm"
)
target_name = target_name.replace(
"pre_feedforward_layernorm", "ff_layer.pre_layer_norm"
)
target_name = target_name.replace(
"post_feedforward_layernorm", "ff_layer.post_layer_norm"
)
target_name = target_name.replace(
"ff_layer.pre_layer_norm.weight", "ff_layer.pre_layer_norm.scale"
)
target_name = target_name.replace(
"ff_layer.post_layer_norm.weight", "ff_layer.post_layer_norm.scale"
)
target_name = target_name.replace(
"post_layer_norm.weight", "post_layer_norm.scale"
)
target_name = target_name.replace("self_attn.q_proj", "self_attention.q")
target_name = target_name.replace("self_attn.k_proj", "self_attention.k")
target_name = target_name.replace("self_attn.v_proj", "self_attention.v")
target_name = target_name.replace("self_attn.o_proj", "self_attention.post")
# Gemma3 QK norms
target_name = target_name.replace(
"self_attn.q_norm", "self_attention.q_norm"
)
target_name = target_name.replace(
"self_attention.q_norm.weight", "self_attention.q_norm.scale"
)
target_name = target_name.replace(
"self_attn.k_norm", "self_attention.k_norm"
)
target_name = target_name.replace(
"self_attention.k_norm.weight", "self_attention.k_norm.scale"
)
target_name = target_name.replace(
"model.embed_tokens", "params.lm.softmax.logits_ffn"
)
# Alternative Gemma3 embedding layer name
target_name = target_name.replace(
"embedder", "params.lm.softmax.logits_ffn"
)
target_name = target_name.replace("model.norm", "params.lm.final_ln")
target_name = target_name.replace("final_ln.weight", "final_ln.scale")
target_name = target_name.replace(".weight", ".w")
# For LoRA weights
if "post" in target_name:
target_name = target_name.replace("lora_A.w", "w_prime_right")
target_name = target_name.replace("lora_B.w", "w_prime_left")
else:
target_name = target_name.replace("lora_A.w", "w_prime_left")
target_name = target_name.replace("lora_B.w", "w_prime_right")
return target_name
class SafetensorsCkptLoader(converter_base.CkptLoaderBase):
"""CkptLoader implementation for loading the Safetensors."""
def __init__(
self,
ckpt_path: str,
is_symmetric: bool,
attention_quant_bits: int,
feedforward_quant_bits: int,
embedding_quant_bits: int,
special_model: str,
backend: str,
):
"""Initializes the loader.
Args:
ckpt_path: The filepath to the safetensors file.
is_symmetric: Whether to apply symmetric or asymmetric quantization.
attention_quant_bits: An integer that specify the target quantization bits
(support 8 or 4) for the attention layers.
feedforward_quant_bits: An integer that specify the target quantization
bits (support 8 or 4) for the feedforward layers in each Transformer
blocks.
embedding_quant_bits: An integer that specify the target quantization bits
(support 8 or 4) for the embedding (and the final projection) layers.
special_model: A string that indicates which input model is and whether
any special treatment is needed.
backend: A string indicating the backend used when converting this model.
Valid options are "cpu" and "gpu".
"""
super().__init__(
ckpt_path,
is_symmetric,
attention_quant_bits,
feedforward_quant_bits,
embedding_quant_bits,
)
self._special_model = special_model
self._reader = _SafetensorsReader(ckpt_path)
if special_model in ["STABLELM_4E1T_3B"]:
self.mapper = StablelmMapper(
is_symmetric,
attention_quant_bits,
feedforward_quant_bits,
embedding_quant_bits,
backend,
self._reader,
)
elif special_model in ["PHI_2"]:
self.mapper = PhiMapper(
is_symmetric,
attention_quant_bits,
feedforward_quant_bits,
embedding_quant_bits,
backend,
self._reader,
)
elif special_model in [
"GEMMA_2B",
"GEMMA_7B",
"GEMMA2_2B",
"GEMMA3_1B",
"GEMMA3_4B",
"GEMMA3_12B",
"GEMMA3_27B",
"GEMMA3_300M",
"FUNCTIONGEMMA3_300M",
]:
self.mapper = GemmaMapper(
is_symmetric,
attention_quant_bits,
feedforward_quant_bits,
embedding_quant_bits,
backend,
self._reader,
False if special_model in ["GEMMA_2B", "GEMMA_7B"] else True,
)
else:
raise ValueError(f"Unknown special model: {special_model}")
def load_to_actions(
self,
) -> Iterator[List[converter_base.QuantizationAction]]:
tensor_names = self._reader.get_tensor_names()
for tensor_name in tensor_names:
tensor_actions = self.mapper.map_to_actions(tensor_name)
if tensor_actions is None:
continue
yield tensor_actions
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/safetensors_converter_test.py
+83
View File
@@ -0,0 +1,83 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Unit tests for safetensors_converter."""
import os
from absl.testing import absltest
from absl.testing import parameterized
from mediapipe.tasks.python.genai.converter import safetensors_converter
from mediapipe.tasks.python.test import test_utils
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
_SAFETENSORS_FILE = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, 'stablelm_3b_4e1t_test_weight.safetensors')
)
class SafetensorsConverterTest(parameterized.TestCase):
VARIABLE_NAMES = [
'model.embed_tokens.weight',
'model.layers.0.input_layernorm.bias',
'model.layers.0.input_layernorm.weight',
'model.layers.0.mlp.down_proj.weight',
'model.layers.0.mlp.gate_proj.weight',
'model.layers.0.mlp.up_proj.weight',
'model.layers.0.post_attention_layernorm.bias',
'model.layers.0.post_attention_layernorm.weight',
'model.layers.0.self_attn.k_proj.weight',
'model.layers.0.self_attn.o_proj.weight',
'model.layers.0.self_attn.q_proj.weight',
'model.layers.0.self_attn.v_proj.weight',
'model.norm.bias',
'model.norm.weight',
'lm_head.weight',
]
def test_init(self):
loader = safetensors_converter.SafetensorsCkptLoader(
ckpt_path=_SAFETENSORS_FILE,
is_symmetric=True,
attention_quant_bits=8,
feedforward_quant_bits=8,
embedding_quant_bits=8,
special_model='STABLELM_4E1T_3B',
backend='gpu',
)
self.assertEqual(loader._ckpt_path, _SAFETENSORS_FILE)
self.assertEqual(loader._is_symmetric, True)
self.assertEqual(loader._attention_quant_bits, 8)
self.assertEqual(loader._feedforward_quant_bits, 8)
@parameterized.product(
quant_bits=(4, 8),
)
def test_load_to_actions(self, quant_bits):
loader = safetensors_converter.SafetensorsCkptLoader(
ckpt_path=_SAFETENSORS_FILE,
is_symmetric=True,
attention_quant_bits=8,
feedforward_quant_bits=quant_bits,
embedding_quant_bits=8,
special_model='STABLELM_4E1T_3B',
backend='gpu',
)
actions = loader.load_to_actions()
self.assertLen(list(actions), 15)
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/weight_bins_writer.py
+127
View File
@@ -0,0 +1,127 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""ModelWriter for writing a set of weights as binary files."""
import contextlib
import os
from typing import Dict, Tuple
import numpy as np
from mediapipe.tasks.python.genai.converter import converter_base
from mediapipe.tasks.python.genai.converter import external_dependencies
from mediapipe.tasks.python.genai.converter import quantization_util
jnp = external_dependencies.jnp
@contextlib.contextmanager
def filemanager(filename: str, mode: str):
try:
with open(filename, mode) as f:
yield f
finally:
pass
def removeprefix(s, prefix):
"""Removes the prefix from a string."""
if s.startswith(prefix):
return s[len(prefix) :]
return s
class WeightBinsWriter(converter_base.ModelWriterBase):
"""A ModelWriter for writing a set of weights as binary files."""
def get_weight_info(self, var_name: str, weight: np.ndarray) -> str:
"""Gets the string that describes the weights."""
dtype_str = str(weight.dtype)
shape_str = '_'.join(map(str, weight.shape))
return f'mdl_vars.{var_name}.{dtype_str}.{shape_str}\n'
def write_variables(
self,
variables: Dict[str, Tuple[np.ndarray, bool]],
use_fake_values: bool = False,
):
"""Writes variable to the binary files. One for each layer.
Args:
variables: A dictionary that maps from the target variable names to the
quantized tensor values along with a boolean that indicates whether to
pack the values (only applicable for the 4-bit quantized tensors).
use_fake_values: Whether to use fake values for the weights. If set to
True, the weights will be filled with zeros.
"""
weights_info = []
for var_name, value in variables.items():
output = value[0]
if use_fake_values:
output.fill(0)
if output.dtype == jnp.int2:
print('saving int2 weight: ', var_name)
output = output.astype(jnp.int8)
output = np.expand_dims(np.ravel(output), axis=-1)
output = quantization_util.pack_2bit_into_int8(output, 0)
if value[1]:
# Squeeze the tensor to make sure it is a 1D array for packing.
output = np.expand_dims(np.ravel(output), axis=-1)
# Extra pack needed for 4 bit. We always pack the weights along the
# first dimension since the tensor has already been squeezed.
output = quantization_util.pack_4bit(output, 0, jnp.int8)
if 'combined_qkv' in var_name:
var_name = removeprefix(var_name, 'mld_vars.')
var_name_q = var_name.replace('combined_qkv', 'q')
var_name_k = var_name.replace('combined_qkv', 'k')
var_name_v = var_name.replace('combined_qkv', 'v')
if output.shape[0] == 3:
weight_q, weight_k, weight_v = output
assert weight_q.shape == weight_k.shape == weight_v.shape
else: # LoRA right weight is shared across q, k, v
weight_q = weight_k = weight_v = output
weights_info.append(self.get_weight_info(var_name_q, weight_q))
path_q = os.path.join(self._output_dir, var_name_q)
with filemanager(path_q, 'wb') as f:
f.write(weight_q.tobytes())
weights_info.append(self.get_weight_info(var_name_k, weight_k))
path_k = os.path.join(self._output_dir, var_name_k)
with filemanager(path_k, 'wb') as f:
f.write(weight_k.tobytes())
path_v = os.path.join(self._output_dir, var_name_v)
with filemanager(path_v, 'wb') as f:
f.write(weight_v.tobytes())
weights_info.append(self.get_weight_info(var_name_v, weight_v))
else:
if 'key' in var_name:
var_name = var_name.replace('key', 'k')
if 'query' in var_name:
var_name = var_name.replace('query', 'q')
if 'value' in var_name:
var_name = var_name.replace('value', 'v')
path = os.path.join(
self._output_dir, removeprefix(var_name, 'mdl_vars.')
)
with filemanager(path, 'wb') as f:
f.write(output.tobytes())
weights_info.append(self.get_weight_info(var_name, output))
# Sort weights_info
weights_info.sort()
with filemanager(
os.path.join(self._output_dir, 'layer_info.txt'), 'a'
) as finfo:
for line in weights_info:
finfo.write(line + '\n')
backend/.venv/Lib/site-packages/mediapipe/tasks/python/genai/converter/weight_bins_writer_test.py
+95
View File
@@ -0,0 +1,95 @@
# Copyright 2024 The MediaPipe Authors.
#
# 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.
"""Unit tests for pax_converter."""
import os
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from mediapipe.tasks.python.genai.converter import weight_bins_writer
class WeightBinsWriterTest(parameterized.TestCase):
def test_get_weight_info(self):
output_dir = os.path.join(flags.FLAGS.test_tmpdir, 'output_dir')
writer = weight_bins_writer.WeightBinsWriter(
output_dir=output_dir, backend='cpu'
)
var_name = 'params.lm.softmax.logits_ffn.linear.w'
weight_info = writer.get_weight_info(
var_name, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
)
self.assertEqual(
weight_info,
'mdl_vars.params.lm.softmax.logits_ffn.linear.w.float32.2_3\n',
)
def test_load_to_actions(self):
output_dir = os.path.join(flags.FLAGS.test_tmpdir, 'output_dir')
writer = weight_bins_writer.WeightBinsWriter(
output_dir=output_dir, backend='cpu'
)
variables = {
'mdl_vars.params.lm.softmax.logits_ffn.linear.w': (
np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
False,
),
}
writer.write_variables(variables)
file_size = os.path.getsize(
os.path.join(output_dir, 'params.lm.softmax.logits_ffn.linear.w')
)
self.assertEqual(file_size, 6 * 4)
@parameterized.named_parameters(
(
'real_weights',
np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
False,
),
(
'fake_weights',
np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
True,
),
)
def test_write_variables(self, var_values, use_fake_values):
output_dir = os.path.join(flags.FLAGS.test_tmpdir, 'output_dir')
writer = weight_bins_writer.WeightBinsWriter(
output_dir=output_dir, backend='gpu'
)
variables = {
'mdl_vars.params.lm.softmax.logits_ffn.linear.w': (
var_values,
False,
),
}
writer.write_variables(variables, use_fake_values=use_fake_values)
with open(
os.path.join(output_dir, 'params.lm.softmax.logits_ffn.linear.w'), 'rb'
) as f:
data = np.frombuffer(f.read(), dtype=np.float32).reshape(var_values.shape)
expected_values = var_values if not use_fake_values else np.zeros_like(
var_values
)
self.assertTrue(np.all(data == expected_values))
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata.py
+1013
View File
File diff suppressed because it is too large Load Diff
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_displayer_cli.py
+34
View File
@@ -0,0 +1,34 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""CLI tool for display metadata."""
from absl import app
from absl import flags
from mediapipe.tasks.python.metadata import metadata
FLAGS = flags.FLAGS
flags.DEFINE_string('model_path', None, 'Path to the TFLite model file.')
flags.DEFINE_string('export_json_path', None, 'Path to the output JSON file.')
def main(_):
displayer = metadata.MetadataDisplayer.with_model_file(FLAGS.model_path)
with open(FLAGS.export_json_path, 'w') as f:
f.write(displayer.get_metadata_json())
if __name__ == '__main__':
app.run(main)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/image_classifier.py
+71
View File
@@ -0,0 +1,71 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Writes metadata and label file to the image classifier models."""
from typing import List, Optional
from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
_MODEL_NAME = "ImageClassifier"
_MODEL_DESCRIPTION = ("Identify the most prominent object in the image from a "
"known set of categories.")
class MetadataWriter(metadata_writer.MetadataWriterBase):
"""MetadataWriter to write the metadata for image classifier."""
@classmethod
def create(
cls,
model_buffer: bytearray,
input_norm_mean: List[float],
input_norm_std: List[float],
labels: metadata_writer.Labels,
score_calibration: Optional[metadata_writer.ScoreCalibration] = None
) -> "MetadataWriter":
"""Creates MetadataWriter to write the metadata for image classifier.
The parameters required in this method are mandatory when using MediaPipe
Tasks.
Note that only the output TFLite is used for deployment. The output JSON
content is used to interpret the metadata content.
Args:
model_buffer: A valid flatbuffer loaded from the TFLite model file.
input_norm_mean: the mean value used in the input tensor normalization
[1].
input_norm_std: the std value used in the input tensor normalizarion [1].
labels: an instance of Labels helper class used in the output
classification tensor [2].
score_calibration: A container of the score calibration operation [3] in
the classification tensor. Optional if the model does not use score
calibration.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
[3]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456
Returns:
A MetadataWriter object.
"""
writer = metadata_writer.MetadataWriter(model_buffer)
writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
writer.add_image_input(input_norm_mean, input_norm_std)
writer.add_classification_output(labels, score_calibration)
return cls(writer)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/image_segmenter.py
+170
View File
@@ -0,0 +1,170 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Writes metadata and label file to the image segmenter models."""
import enum
from typing import List, Optional
import flatbuffers
from mediapipe.tasks.metadata import image_segmenter_metadata_schema_py_generated as _segmenter_metadata_fb
from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb
from mediapipe.tasks.python.metadata import metadata
from mediapipe.tasks.python.metadata.metadata_writers import metadata_info
from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
from mediapipe.tasks.python.metadata.metadata_writers import writer_utils
_MODEL_NAME = "ImageSegmenter"
_MODEL_DESCRIPTION = (
"Semantic image segmentation predicts whether each pixel "
"of an image is associated with a certain class."
)
# Metadata Schema file for image segmenter.
_FLATC_METADATA_SCHEMA_FILE = metadata.get_path_to_datafile(
"../../../metadata/image_segmenter_metadata_schema.fbs",
)
# Metadata name in custom metadata field. The metadata name is used to get
# image segmenter metadata from SubGraphMetadata.custom_metadata and
# shouldn't be changed.
_METADATA_NAME = "SEGMENTER_METADATA"
class Activation(enum.Enum):
NONE = 0
SIGMOID = 1
SOFTMAX = 2
# Create an individual method for getting the metadata json file, so that it can
# be used as a standalone util.
def convert_to_json(metadata_buffer: bytearray) -> str:
"""Converts the metadata into a json string.
Args:
metadata_buffer: valid metadata buffer in bytes.
Returns:
Metadata in JSON format.
Raises:
ValueError: error occurred when parsing the metadata schema file.
"""
return metadata.convert_to_json(
metadata_buffer,
custom_metadata_schema={_METADATA_NAME: _FLATC_METADATA_SCHEMA_FILE},
)
class ImageSegmenterOptionsMd(metadata_info.CustomMetadataMd):
"""Image segmenter options metadata."""
_METADATA_FILE_IDENTIFIER = b"V001"
def __init__(self, activation: Activation) -> None:
"""Creates an ImageSegmenterOptionsMd object.
Args:
activation: activation function of the output layer in the image
segmenter.
"""
self.activation = activation
super().__init__(name=_METADATA_NAME)
def create_metadata(self) -> _metadata_fb.CustomMetadataT:
"""Creates the image segmenter options metadata.
Returns:
A Flatbuffers Python object of the custom metadata including image
segmenter options metadata.
"""
segmenter_options = _segmenter_metadata_fb.ImageSegmenterOptionsT()
segmenter_options.activation = self.activation.value
# Get the image segmenter options flatbuffer.
b = flatbuffers.Builder(0)
b.Finish(segmenter_options.Pack(b), self._METADATA_FILE_IDENTIFIER)
segmenter_options_buf = b.Output()
# Add the image segmenter options flatbuffer in custom metadata.
custom_metadata = _metadata_fb.CustomMetadataT()
custom_metadata.name = self.name
custom_metadata.data = segmenter_options_buf
return custom_metadata
class MetadataWriter(metadata_writer.MetadataWriterBase):
"""MetadataWriter to write the metadata for image segmenter."""
@classmethod
def create(
cls,
model_buffer: bytearray,
input_norm_mean: List[float],
input_norm_std: List[float],
labels: Optional[metadata_writer.Labels] = None,
activation: Optional[Activation] = None,
) -> "MetadataWriter":
"""Creates MetadataWriter to write the metadata for image segmenter.
The parameters required in this method are mandatory when using MediaPipe
Tasks.
Example usage:
metadata_writer = image_segmenter.Metadatawriter.create(model_buffer, ...)
tflite_content, json_content = metadata_writer.populate()
When calling `populate` function in this class, it returns TfLite content
and JSON content. Note that only the output TFLite is used for deployment.
The output JSON content is used to interpret the metadata content.
Args:
model_buffer: A valid flatbuffer loaded from the TFLite model file.
input_norm_mean: the mean value used in the input tensor normalization
[1].
input_norm_std: the std value used in the input tensor normalizarion [1].
labels: an instance of Labels helper class used in the output category
tensor [2].
activation: activation function for the output layer.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L116
Returns:
A MetadataWriter object.
"""
writer = metadata_writer.MetadataWriter(model_buffer)
writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
writer.add_image_input(input_norm_mean, input_norm_std)
if activation is not None:
option_md = ImageSegmenterOptionsMd(activation)
writer.add_custom_metadata(option_md)
num_output_tensors = writer_utils.get_subgraph(model_buffer).OutputsLength()
if num_output_tensors == 2:
# For image segmenter model with 2 output tensors, the first one is
# quality score, and the second one is matting mask.
writer.add_feature_output(
"quality score", "The quality score of matting result."
)
writer.add_segmentation_output(labels=labels)
return cls(writer)
def populate(self) -> tuple[bytearray, str]:
model_buf, _ = super().populate()
metadata_buf = metadata.get_metadata_buffer(model_buf)
json_content = convert_to_json(metadata_buf)
return model_buf, json_content
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py
+1166
View File
File diff suppressed because it is too large Load Diff
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py
+845
View File
@@ -0,0 +1,845 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Generic metadata writer."""
import collections
import csv
import dataclasses
import os
import tempfile
from typing import List, Optional, Tuple, Union
import flatbuffers
from mediapipe.tasks.metadata import metadata_schema_py_generated as metadata_fb
from mediapipe.tasks.python.metadata import metadata
from mediapipe.tasks.python.metadata.metadata_writers import metadata_info
from mediapipe.tasks.python.metadata.metadata_writers import writer_utils
_INPUT_IMAGE_NAME = 'image'
_INPUT_IMAGE_DESCRIPTION = 'Input image to be processed.'
_INPUT_REGEX_TEXT_NAME = 'input_text'
_INPUT_REGEX_TEXT_DESCRIPTION = ('Embedding vectors representing the input '
'text to be processed.')
_OUTPUT_CLASSIFICATION_NAME = 'score'
_OUTPUT_CLASSIFICATION_DESCRIPTION = 'Score of the labels respectively.'
_OUTPUT_SEGMENTATION_MASKS_NAME = 'segmentation_masks'
_OUTPUT_SEGMENTATION_MASKS_DESCRIPTION = (
'Masks over the target objects with high accuracy.'
)
# Detection tensor result to be grouped together.
_DETECTION_GROUP_NAME = 'detection_result'
# File name to export score calibration parameters.
_SCORE_CALIBATION_FILENAME = 'score_calibration.txt'
@dataclasses.dataclass
class CalibrationParameter:
"""Parameters for score calibration [1].
Score calibration is performed on an output tensor through sigmoid functions.
One of the main purposes of score calibration is to make scores across classes
comparable, so that a common threshold can be used for all output classes.
For each index in the output tensor, this applies:
* `f(x) = scale / (1 + e^-(slope * g(x) + offset))` if `x > min_score` or if
no `min_score` has been specified.
* `f(x) = default_score` otherwise or if no scale, slope and offset have
been specified.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L434
"""
scale: float
slope: float
offset: float
min_score: Optional[float] = None
@dataclasses.dataclass
class LabelItem:
"""Label item for labels per locale.
Attributes:
filename: The file name to save the labels.
names: A list of label names.
locale: The specified locale for labels.
"""
filename: str
names: List[str]
locale: Optional[str] = None
@dataclasses.dataclass
class ScoreThresholding:
"""Parameters to performs thresholding on output tensor values [1].
Attributes:
global_score_threshold: The recommended global threshold below which results
are considered low-confidence and should be filtered out. [1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468
"""
global_score_threshold: float
@dataclasses.dataclass
class RegexTokenizer:
"""Parameters of the Regex tokenizer [1] metadata information.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
Attributes:
delim_regex_pattern: the regular expression to segment strings and create
tokens.
vocab_file_path: path to the vocabulary file.
"""
delim_regex_pattern: str
vocab_file_path: str
@dataclasses.dataclass
class BertTokenizer:
"""Parameters of the Bert tokenizer [1] metadata information.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
Attributes:
vocab_file_path: path to the vocabulary file.
"""
vocab_file_path: str
@dataclasses.dataclass
class SentencePieceTokenizer:
"""Parameters of the sentence piece tokenizer tokenizer [1] metadata information.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
Attributes:
sentence_piece_model_path: path to the sentence piece model file.
vocab_file_path: path to the vocabulary file.
"""
sentence_piece_model_path: str
vocab_file_path: Optional[str] = None
class Labels(object):
"""Simple container holding classification labels of a particular tensor.
Example usage:
# The first added label list can be used as category names as needed.
labels = Labels()
.add(['/m/011l78', '/m/031d23'])
.add(['cat', 'dog], 'en')
.add(['chat', 'chien], 'fr')
"""
def __init__(self) -> None:
self._labels = [] # [LabelItem]
@property
def labels(self) -> List[LabelItem]:
return self._labels
def add(self,
labels: List[str],
locale: Optional[str] = None,
exported_filename: Optional[str] = None) -> 'Labels':
"""Adds labels in the container.
Args:
labels: A list of label names, e.g. ['apple', 'pear', 'banana'].
locale: The specified locale for labels.
exported_filename: The file name to export the labels. If not set,
filename defaults to 'labels.txt'.
Returns:
The Labels instance, can be used for chained operation.
"""
if not labels:
raise ValueError('The list of labels is empty.')
# Prepare the new item to be inserted
if not exported_filename:
exported_filename = 'labels'
if locale:
exported_filename += f'_{locale}'
exported_filename += '.txt'
item = LabelItem(filename=exported_filename, names=labels, locale=locale)
# Insert the new element at the end of the list
self._labels.append(item)
return self
def add_from_file(self,
label_filepath: str,
locale: Optional[str] = None,
exported_filename: Optional[str] = None) -> 'Labels':
"""Adds a label file in the container.
Args:
label_filepath: File path to read labels. Each line is a label name in the
file.
locale: The specified locale for labels.
exported_filename: The file name to export the labels. If not set,
filename defaults to 'labels.txt'.
Returns:
The Labels instance, can be used for chained operation.
"""
with open(label_filepath, 'r') as f:
labels = f.read().split('\n')
return self.add(labels, locale, exported_filename)
class ScoreCalibration:
"""Simple container holding score calibration related parameters."""
# A shortcut to avoid client side code importing metadata_fb
transformation_types = metadata_fb.ScoreTransformationType
def __init__(self,
transformation_type: metadata_fb.ScoreTransformationType,
parameters: List[Optional[CalibrationParameter]],
default_score: int = 0):
self.transformation_type = transformation_type
self.parameters = parameters
self.default_score = default_score
@classmethod
def create_from_file(cls,
transformation_type: metadata_fb.ScoreTransformationType,
file_path: str,
default_score: int = 0) -> 'ScoreCalibration':
"""Creates ScoreCalibration from the file.
Args:
transformation_type: type of the function used for transforming the
uncalibrated score before applying score calibration.
file_path: file_path of the score calibration file [1]. Contains
sigmoid-based score calibration parameters, formatted as CSV. Lines
contain for each index of an output tensor the scale, slope, offset and
(optional) min_score parameters to be used for sigmoid fitting (in this
order and in `strtof`-compatible [2] format). Scale should be a
non-negative value. A line may be left empty to default calibrated
scores for this index to default_score. In summary, each line should
thus contain 0, 3 or 4 comma-separated values.
default_score: the default calibrated score to apply if the uncalibrated
score is below min_score or if no parameters were specified for a given
index.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L133
[2]:
https://en.cppreference.com/w/c/string/byte/strtof
Returns:
A ScoreCalibration object.
Raises:
ValueError: if the score_calibration file is malformed.
"""
with open(file_path, 'r') as calibration_file:
csv_reader = csv.reader(calibration_file, delimiter=',')
parameters = []
for row in csv_reader:
if not row:
parameters.append(None)
continue
if len(row) != 3 and len(row) != 4:
raise ValueError(
f'Expected empty lines or 3 or 4 parameters per line in score'
f' calibration file, but got {len(row)}.')
if float(row[0]) < 0:
raise ValueError(
f'Expected scale to be a non-negative value, but got '
f'{float(row[0])}.')
parameters.append(
CalibrationParameter(
scale=float(row[0]),
slope=float(row[1]),
offset=float(row[2]),
min_score=None if len(row) == 3 else float(row[3])))
return cls(transformation_type, parameters, default_score)
def _fill_default_tensor_names(
tensor_metadata_list: List[metadata_fb.TensorMetadataT],
tensor_names_from_model: List[str]):
"""Fills the default tensor names."""
# If tensor name in metadata is empty, default to the tensor name saved in
# the model.
for tensor_metadata, name in zip(tensor_metadata_list,
tensor_names_from_model):
tensor_metadata.name = tensor_metadata.name or name
def _pair_tensor_metadata(
tensor_md: List[metadata_info.TensorMd],
tensor_names_from_model: List[str]) -> List[metadata_info.TensorMd]:
"""Pairs tensor_md according to the tensor names from the model."""
tensor_names_from_arg = [
md.tensor_name for md in tensor_md or [] if md.tensor_name is not None
]
if not tensor_names_from_arg:
return tensor_md
if collections.Counter(tensor_names_from_arg) != collections.Counter(
tensor_names_from_model):
raise ValueError(
'The tensor names from arguments ({}) do not match the tensor names'
' read from the model ({}).'.format(tensor_names_from_arg,
tensor_names_from_model))
pairs_tensor_md = []
name_md_dict = dict(zip(tensor_names_from_arg, tensor_md))
for name in tensor_names_from_model:
pairs_tensor_md.append(name_md_dict[name])
return pairs_tensor_md
def _create_metadata_buffer(
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[List[metadata_info.TensorMd]] = None,
output_md: Optional[List[metadata_info.TensorMd]] = None,
input_process_units: Optional[List[metadata_fb.ProcessUnitT]] = None,
output_group_md: Optional[List[metadata_info.TensorGroupMd]] = None,
custom_metadata_md: Optional[List[metadata_info.CustomMetadataMd]] = None,
) -> bytearray:
"""Creates a buffer of the metadata.
Args:
model_buffer: valid buffer of the model file.
general_md: general information about the model.
input_md: metadata information of the input tensors.
output_md: metadata information of the output tensors.
input_process_units: a lists of metadata of the input process units [1].
output_group_md: a list of metadata of output tensor groups [2];
custom_metadata_md: a lists of custom metadata.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L655
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L677
Returns:
A buffer of the metadata.
Raises:
ValueError: if the tensor names from `input_md` and `output_md` do not
match the tensor names read from the model.
"""
# Create input metadata from `input_md`.
if input_md:
input_md = _pair_tensor_metadata(
input_md, writer_utils.get_input_tensor_names(model_buffer))
input_metadata = [m.create_metadata() for m in input_md]
else:
num_input_tensors = writer_utils.get_subgraph(model_buffer).InputsLength()
input_metadata = [metadata_fb.TensorMetadataT()] * num_input_tensors
_fill_default_tensor_names(input_metadata,
writer_utils.get_input_tensor_names(model_buffer))
# Create output metadata from `output_md`.
if output_md:
output_md = _pair_tensor_metadata(
output_md, writer_utils.get_output_tensor_names(model_buffer))
output_metadata = [m.create_metadata() for m in output_md]
else:
num_output_tensors = writer_utils.get_subgraph(model_buffer).OutputsLength()
output_metadata = [metadata_fb.TensorMetadataT()] * num_output_tensors
_fill_default_tensor_names(output_metadata,
writer_utils.get_output_tensor_names(model_buffer))
# Create the subgraph metadata.
subgraph_metadata = metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = input_metadata
subgraph_metadata.outputTensorMetadata = output_metadata
if input_process_units:
subgraph_metadata.inputProcessUnits = input_process_units
if custom_metadata_md:
subgraph_metadata.customMetadata = [
m.create_metadata() for m in custom_metadata_md
]
if output_group_md:
subgraph_metadata.outputTensorGroups = [
m.create_metadata() for m in output_group_md
]
# Create the whole model metadata.
if general_md is None:
general_md = metadata_info.GeneralMd()
model_metadata = general_md.create_metadata()
model_metadata.subgraphMetadata = [subgraph_metadata]
# Get the metadata flatbuffer.
b = flatbuffers.Builder(0)
b.Finish(
model_metadata.Pack(b),
metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
return b.Output()
class MetadataWriter(object):
"""Generic Metadata writer.
Example usage:
For an example model which requires two inputs: image and general feature
inputs, and generates one output: classification.
with open(model_path, 'rb') as f:
writer = MetadataWriter.create(f.read())
model_content, metadata_json_content = writer
.add_genernal_info('model_name', 'model description')
.add_image_input()
.add_feature_input()
.add_classification_output(Labels().add(['A', 'B']))
.populate()
"""
@classmethod
def create(cls, model_buffer: bytearray) -> 'MetadataWriter':
return cls(model_buffer)
def __init__(self, model_buffer: bytearray) -> None:
self._model_buffer = model_buffer
self._general_md = None
self._input_mds = []
self._input_process_units = []
self._output_mds = []
self._output_group_mds = []
self._associated_files = []
self._custom_metadata_mds = []
self._temp_folder = tempfile.TemporaryDirectory()
def __del__(self):
if os.path.exists(self._temp_folder.name):
self._temp_folder.cleanup()
def add_general_info(
self,
model_name: str,
model_description: Optional[str] = None) -> 'MetadataWriter':
"""Adds a general info metadata for the general metadata informantion."""
# Will overwrite the previous `self._general_md` if exists.
self._general_md = metadata_info.GeneralMd(
name=model_name, description=model_description)
return self
color_space_types = metadata_fb.ColorSpaceType
def add_feature_input(self,
name: Optional[str] = None,
description: Optional[str] = None) -> 'MetadataWriter':
"""Adds an input tensor metadata for the general basic feature input."""
input_md = metadata_info.TensorMd(name=name, description=description)
self._input_mds.append(input_md)
return self
def add_image_input(
self,
norm_mean: List[float],
norm_std: List[float],
color_space_type: Optional[int] = metadata_fb.ColorSpaceType.RGB,
name: str = _INPUT_IMAGE_NAME,
description: str = _INPUT_IMAGE_DESCRIPTION) -> 'MetadataWriter':
"""Adds an input image metadata for the image input.
Args:
norm_mean: The mean value used to normalize each input channel. If there
is only one element in the list, its value will be broadcasted to all
channels. Also note that norm_mean and norm_std should have the same
number of elements. [1]
norm_std: The std value used to normalize each input channel. If there is
only one element in the list, its value will be broadcasted to all
channels. [1]
color_space_type: The color space type of the input image. [2]
name: Name of the input tensor.
description: Description of the input tensor.
Returns:
The MetadataWriter instance, can be used for chained operation.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L198
"""
input_md = metadata_info.InputImageTensorMd(
name=name,
description=description,
norm_mean=norm_mean,
norm_std=norm_std,
color_space_type=color_space_type,
tensor_type=self._input_tensor_type(len(self._input_mds)))
self._input_mds.append(input_md)
return self
def add_regex_text_input(
self,
regex_tokenizer: RegexTokenizer,
name: str = _INPUT_REGEX_TEXT_NAME,
description: str = _INPUT_REGEX_TEXT_DESCRIPTION) -> 'MetadataWriter':
"""Adds an input text metadata for the text input with regex tokenizer.
Args:
regex_tokenizer: information of the regex tokenizer [1] used to process
the input string.
name: Name of the input tensor.
description: Description of the input tensor.
Returns:
The MetadataWriter instance, can be used for chained operation.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
"""
tokenizer_md = metadata_info.RegexTokenizerMd(
delim_regex_pattern=regex_tokenizer.delim_regex_pattern,
vocab_file_path=regex_tokenizer.vocab_file_path)
input_md = metadata_info.InputTextTensorMd(
name=name, description=description, tokenizer_md=tokenizer_md)
self._input_mds.append(input_md)
self._associated_files.append(regex_tokenizer.vocab_file_path)
return self
def add_bert_text_input(self, tokenizer: Union[BertTokenizer,
SentencePieceTokenizer],
ids_name: str, mask_name: str,
segment_name: str) -> 'MetadataWriter':
"""Adds an metadata for the text input with bert / sentencepiece tokenizer.
`ids_name`, `mask_name`, and `segment_name` correspond to the `Tensor.name`
in the TFLite schema, which help to determine the tensor order when
populating metadata.
Args:
tokenizer: information of the tokenizer used to process the input string,
if any. Supported tokenziers are: `BertTokenizer` [1] and
`SentencePieceTokenizer` [2].
ids_name: name of the ids tensor, which represents the tokenized ids of
the input text.
mask_name: name of the mask tensor, which represents the mask with `1` for
real tokens and `0` for padding tokens.
segment_name: name of the segment ids tensor, where `0` stands for the
first sequence, and `1` stands for the second sequence if exists.
Returns:
The MetadataWriter instance, can be used for chained operation.
Raises:
ValueError: if the type tokenizer is not BertTokenizer or
SentencePieceTokenizer.
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
"""
if isinstance(tokenizer, BertTokenizer):
tokenizer_md = metadata_info.BertTokenizerMd(
vocab_file_path=tokenizer.vocab_file_path)
elif isinstance(tokenizer, SentencePieceTokenizer):
tokenizer_md = metadata_info.SentencePieceTokenizerMd(
sentence_piece_model_path=tokenizer.sentence_piece_model_path,
vocab_file_path=tokenizer.vocab_file_path)
else:
raise ValueError(
f'The type of tokenizer, {type(tokenizer)}, is unsupported')
bert_input_md = metadata_info.BertInputTensorsMd(
self._model_buffer,
ids_name,
mask_name,
segment_name,
tokenizer_md=tokenizer_md)
self._input_mds.extend(bert_input_md.input_md)
self._associated_files.extend(
bert_input_md.get_tokenizer_associated_files())
self._input_process_units.extend(
bert_input_md.create_input_process_unit_metadata())
return self
def add_classification_output(
self,
labels: Optional[Labels] = None,
score_calibration: Optional[ScoreCalibration] = None,
score_thresholding: Optional[ScoreThresholding] = None,
name: str = _OUTPUT_CLASSIFICATION_NAME,
description: str = _OUTPUT_CLASSIFICATION_DESCRIPTION
) -> 'MetadataWriter':
"""Add a classification head metadata for classification output tensor.
Example usage:
writer.add_classification_output(
Labels()
.add(['/m/011l78', '/m/031d23'])
.add(['cat', 'dog], 'en')
.add(['chat', 'chien], 'fr')
)
Args:
labels: an instance of Labels helper class.
score_calibration: an instance of ScoreCalibration helper class.
score_thresholding: an instance of ScoreThresholding.
name: Metadata name of the tensor. Note that this is different from tensor
name in the flatbuffer.
description: human readable description of what the output is.
Returns:
The current Writer instance to allow chained operation.
"""
calibration_md = self._create_score_calibration_md(score_calibration)
score_thresholding_md = None
if score_thresholding:
score_thresholding_md = metadata_info.ScoreThresholdingMd(
score_thresholding.global_score_threshold)
label_files = self._create_label_file_md(labels)
output_md = metadata_info.ClassificationTensorMd(
name=name,
description=description,
label_files=label_files,
tensor_type=self._output_tensor_type(len(self._output_mds)),
score_calibration_md=calibration_md,
score_thresholding_md=score_thresholding_md,
)
self._output_mds.append(output_md)
return self
def add_detection_output(
self,
labels: Optional[Labels] = None,
score_calibration: Optional[ScoreCalibration] = None,
group_name: str = _DETECTION_GROUP_NAME,
) -> 'MetadataWriter':
"""Adds a detection head metadata for detection output tensor of models with postprocessing.
Args:
labels: an instance of Labels helper class.
score_calibration: an instance of ScoreCalibration helper class.
group_name: name of output tensor group.
Returns:
The current Writer instance to allow chained operation.
"""
calibration_md = self._create_score_calibration_md(score_calibration)
label_files = self._create_label_file_md(labels)
detection_output_mds = metadata_info.DetectionOutputTensorsMd(
self._model_buffer,
label_files=label_files,
score_calibration_md=calibration_md,
).output_mds
self._output_mds.extend(detection_output_mds)
# Outputs are location, category, score, number of detections.
if len(detection_output_mds) != 4:
raise ValueError('The size of detections output should be 4.')
# The first 3 tensors (location, category, score) are grouped.
group_md = metadata_info.TensorGroupMd(
name=group_name,
tensor_names=[output_md.name for output_md in detection_output_mds[:3]],
)
self._output_group_mds.append(group_md)
return self
def add_raw_detection_output(
self,
labels: Optional[Labels] = None,
output_tensors_order: metadata_info.RawDetectionOutputTensorsOrder = metadata_info.RawDetectionOutputTensorsOrder.UNSPECIFIED,
) -> 'MetadataWriter':
"""Adds a detection head metadata for detection output tensor of models without postprocessing.
Args:
labels: an instance of Labels helper class.
output_tensors_order: the order of the output tensors. For models of
out-of-graph non-maximum-suppression only.
Returns:
The current Writer instance to allow chained operation.
"""
label_files = self._create_label_file_md(labels)
detection_output_mds = metadata_info.RawDetectionOutputTensorsMd(
self._model_buffer,
label_files=label_files,
output_tensors_order=output_tensors_order,
).output_mds
self._output_mds.extend(detection_output_mds)
# Outputs are location, score.
if len(detection_output_mds) != 2:
raise ValueError('The size of detections output should be 2.')
return self
def add_segmentation_output(
self,
labels: Optional[Labels] = None,
name: str = _OUTPUT_SEGMENTATION_MASKS_NAME,
description: str = _OUTPUT_SEGMENTATION_MASKS_DESCRIPTION,
) -> 'MetadataWriter':
"""Adds a segmentation head metadata for segmentation output tensor.
Args:
labels: an instance of Labels helper class.
name: Metadata name of the tensor. Note that this is different from tensor
name in the flatbuffer.
description: human readable description of what the output is.
Returns:
The current Writer instance to allow chained operation.
"""
label_files = self._create_label_file_md(labels)
output_md = metadata_info.SegmentationMaskMd(
name=name,
description=description,
label_files=label_files,
)
self._output_mds.append(output_md)
return self
def add_feature_output(self,
name: Optional[str] = None,
description: Optional[str] = None) -> 'MetadataWriter':
"""Adds an output tensor metadata for the general basic feature output."""
output_md = metadata_info.TensorMd(name=name, description=description)
self._output_mds.append(output_md)
return self
def add_custom_metadata(
self, custom_metadata_md: metadata_info.CustomMetadataMd
) -> 'MetadataWriter':
self._custom_metadata_mds.append(custom_metadata_md)
return self
def populate(self) -> Tuple[bytearray, str]:
"""Populates metadata into the TFLite file.
Note that only the output tflite is used for deployment. The output JSON
content is used to interpret the metadata content.
Returns:
A tuple of (model_with_metadata_in_bytes, metadata_json_content)
"""
# Populates metadata and associated files into TFLite model buffer.
populator = metadata.MetadataPopulator.with_model_buffer(self._model_buffer)
metadata_buffer = _create_metadata_buffer(
model_buffer=self._model_buffer,
general_md=self._general_md,
input_md=self._input_mds,
output_md=self._output_mds,
input_process_units=self._input_process_units,
custom_metadata_md=self._custom_metadata_mds,
output_group_md=self._output_group_mds,
)
populator.load_metadata_buffer(metadata_buffer)
if self._associated_files:
populator.load_associated_files(self._associated_files)
populator.populate()
tflite_content = populator.get_model_buffer()
displayer = metadata.MetadataDisplayer.with_model_buffer(tflite_content)
metadata_json_content = displayer.get_metadata_json()
return tflite_content, metadata_json_content
def _input_tensor_type(self, idx):
return writer_utils.get_input_tensor_types(self._model_buffer)[idx]
def _output_tensor_type(self, idx):
return writer_utils.get_output_tensor_types(self._model_buffer)[idx]
def _export_labels(self, filename: str, index_to_label: List[str]) -> str:
filepath = os.path.join(self._temp_folder.name, filename)
with open(filepath, 'w') as f:
f.write('\n'.join(index_to_label))
self._associated_files.append(filepath)
return filepath
def _export_calibration_file(self, filename: str,
calibrations: List[CalibrationParameter]) -> str:
"""Stores calibration parameters in a csv file."""
filepath = os.path.join(self._temp_folder.name, filename)
with open(filepath, 'w') as f:
for item in calibrations:
if item:
if item.scale is None or item.slope is None or item.offset is None:
raise ValueError('scale, slope and offset values can not be set to '
'None.')
elif item.min_score is not None:
f.write(f'{item.scale},{item.slope},{item.offset},{item.min_score}')
else:
f.write(f'{item.scale},{item.slope},{item.offset}')
f.write('\n')
self._associated_files.append(filepath)
return filepath
def _create_score_calibration_md(
self, score_calibration: ScoreCalibration
) -> Optional[metadata_info.ScoreCalibrationMd]:
"""Creates the ScoreCalibrationMd object."""
if score_calibration is None:
return None
return metadata_info.ScoreCalibrationMd(
score_transformation_type=score_calibration.transformation_type,
default_score=score_calibration.default_score,
file_path=self._export_calibration_file(
_SCORE_CALIBATION_FILENAME, score_calibration.parameters
),
)
def _create_label_file_md(
self, labels: Optional[Labels] = None
) -> Optional[List[metadata_info.LabelFileMd]]:
"""Creates a list of LabelFileMd objects."""
label_files = None
if labels:
label_files = []
for item in labels.labels:
label_files.append(
metadata_info.LabelFileMd(
self._export_labels(item.filename, item.names),
locale=item.locale,
)
)
return label_files
class MetadataWriterBase:
"""Base MetadataWriter class which contains the apis exposed to users.
MetadataWriter for Tasks e.g. image classifier / object detector will inherit
this class for their own usage.
"""
def __init__(self, writer: MetadataWriter) -> None:
self.writer = writer
def populate(self) -> Tuple[bytearray, str]:
"""Populates metadata into the TFLite file.
Note that only the output tflite is used for deployment. The output JSON
content is used to interpret the metadata content.
Returns:
A tuple of (model_with_metadata_in_bytes, metadata_json_content)
"""
return self.writer.populate()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/model_asset_bundle_utils.py
+71
View File
@@ -0,0 +1,71 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Utility methods for creating the model asset bundles."""
from typing import Dict
import zipfile
# Alignment that ensures that all uncompressed files in the model bundle file
# are aligned relative to the start of the file. This lets the files be
# accessed directly via mmap.
_ALIGNMENT = 4
class AlignZipFile(zipfile.ZipFile):
"""ZipFile that stores uncompressed files at particular alignment."""
def __init__(self, *args, alignment: int, **kwargs) -> None:
super().__init__(*args, **kwargs)
assert alignment > 0
self._alignment = alignment
def _writecheck(self, zinfo: zipfile.ZipInfo) -> None:
# Aligned the uncompressed files.
if zinfo.compress_type == zipfile.ZIP_STORED:
offset = self.fp.tell()
header_length = len(zinfo.FileHeader())
padding_length = (
self._alignment - (offset + header_length) % self._alignment
)
if padding_length:
offset += padding_length
self.fp.write(b"\x00" * padding_length)
assert self.fp.tell() == offset
zinfo.header_offset = offset
else:
raise ValueError(
"Only support the uncompressed file (compress_type =="
" zipfile.ZIP_STORED) in zip. The current file compress type is "
+ str(zinfo.compress_type)
)
super()._writecheck(zinfo)
def create_model_asset_bundle(
input_models: Dict[str, bytes], output_path: str
) -> None:
"""Creates the model asset bundle.
Args:
input_models: A dict of input models with key as the model file name and
value as the model content.
output_path: The output file path to save the model asset bundle.
"""
if not input_models or len(input_models) < 2:
raise ValueError("Needs at least two input models for model asset bundle.")
with AlignZipFile(output_path, mode="w", alignment=_ALIGNMENT) as zf:
for file_name, file_buffer in input_models.items():
zf.writestr(file_name, file_buffer)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/object_detector.py
+331
View File
@@ -0,0 +1,331 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Writes metadata and label file to the Object Detector models."""
import dataclasses
from typing import List, Optional
import flatbuffers
from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb
from mediapipe.tasks.metadata import object_detector_metadata_schema_py_generated as _detector_metadata_fb
from mediapipe.tasks.python.metadata import metadata
from mediapipe.tasks.python.metadata.metadata_writers import metadata_info
from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
_MODEL_NAME = "ObjectDetector"
_MODEL_DESCRIPTION = (
"Identify which of a known set of objects might be present and provide "
"information about their positions within the given image or a video "
"stream."
)
# Metadata Schema file for object detector.
_FLATC_METADATA_SCHEMA_FILE = metadata.get_path_to_datafile(
"../../../metadata/object_detector_metadata_schema.fbs",
)
# Metadata name in custom metadata field. The metadata name is used to get
# object detector metadata from SubGraphMetadata.custom_metadata and shouldn't
# be changed.
_METADATA_NAME = "DETECTOR_METADATA"
@dataclasses.dataclass
class FixedAnchor:
"""A fixed size anchor."""
x_center: float
y_center: float
width: Optional[float]
height: Optional[float]
@dataclasses.dataclass
class FixedAnchorsSchema:
"""The schema for a list of anchors with fixed size."""
anchors: List[FixedAnchor]
@dataclasses.dataclass
class SsdAnchorsOptions:
"""The ssd anchors options used in object detector model."""
fixed_anchors_schema: Optional[FixedAnchorsSchema]
@dataclasses.dataclass
class TensorsDecodingOptions:
"""The decoding options to convert model output tensors to detections."""
# The number of output classes predicted by the detection model.
num_classes: int
# The number of output boxes predicted by the detection model.
num_boxes: int
# The number of output values per boxes predicted by the detection
# model. The values contain bounding boxes, keypoints, etc.
num_coords: int
# The offset of keypoint coordinates in the location tensor.
keypoint_coord_offset: int
# The number of predicted keypoints.
num_keypoints: int
# The dimension of each keypoint, e.g. number of values predicted for each
# keypoint.
num_values_per_keypoint: int
# Parameters for decoding SSD detection model.
x_scale: float
y_scale: float
w_scale: float
h_scale: float
# Whether to apply exponential on box size.
apply_exponential_on_box_size: bool
# Whether to apply sigmod function on the score.
sigmoid_score: bool
# Create an individual method for getting the metadata json file, so that it can
# be used as a standalone util.
def convert_to_json(metadata_buffer: bytearray) -> str:
"""Converts the metadata into a json string.
Args:
metadata_buffer: valid metadata buffer in bytes.
Returns:
Metadata in JSON format.
Raises:
ValueError: error occurred when parsing the metadata schema file.
"""
return metadata.convert_to_json(
metadata_buffer,
custom_metadata_schema={_METADATA_NAME: _FLATC_METADATA_SCHEMA_FILE},
)
class ObjectDetectorOptionsMd(metadata_info.CustomMetadataMd):
"""Object detector options metadata."""
_METADATA_FILE_IDENTIFIER = b"V001"
def __init__(
self,
ssd_anchors_options: SsdAnchorsOptions,
tensors_decoding_options: TensorsDecodingOptions,
) -> None:
"""Creates an ObjectDetectorOptionsMd object.
Args:
ssd_anchors_options: the ssd anchors options associated to the object
detector model.
tensors_decoding_options: the tensors decoding options used to decode the
object detector model output.
"""
if ssd_anchors_options.fixed_anchors_schema is None:
raise ValueError(
"Currently only support FixedAnchorsSchema, which cannot be found"
" in ssd_anchors_options."
)
self.ssd_anchors_options = ssd_anchors_options
self.tensors_decoding_options = tensors_decoding_options
super().__init__(name=_METADATA_NAME)
def create_metadata(self) -> _metadata_fb.CustomMetadataT:
"""Creates the image segmenter options metadata.
Returns:
A Flatbuffers Python object of the custom metadata including object
detector options metadata.
"""
detector_options = _detector_metadata_fb.ObjectDetectorOptionsT()
# Set ssd_anchors_options.
ssd_anchors_options = _detector_metadata_fb.SsdAnchorsOptionsT()
fixed_anchors_schema = _detector_metadata_fb.FixedAnchorsSchemaT()
fixed_anchors_schema.anchors = []
for anchor in self.ssd_anchors_options.fixed_anchors_schema.anchors:
anchor_t = _detector_metadata_fb.FixedAnchorT()
anchor_t.xCenter = anchor.x_center
anchor_t.yCenter = anchor.y_center
anchor_t.width = anchor.width
anchor_t.height = anchor.height
fixed_anchors_schema.anchors.append(anchor_t)
ssd_anchors_options.fixedAnchorsSchema = fixed_anchors_schema
detector_options.ssdAnchorsOptions = ssd_anchors_options
# Set tensors_decoding_options.
tensors_decoding_options = _detector_metadata_fb.TensorsDecodingOptionsT()
tensors_decoding_options.numClasses = (
self.tensors_decoding_options.num_classes
)
tensors_decoding_options.numBoxes = self.tensors_decoding_options.num_boxes
tensors_decoding_options.numCoords = (
self.tensors_decoding_options.num_coords
)
tensors_decoding_options.keypointCoordOffset = (
self.tensors_decoding_options.keypoint_coord_offset
)
tensors_decoding_options.numKeypoints = (
self.tensors_decoding_options.num_keypoints
)
tensors_decoding_options.numValuesPerKeypoint = (
self.tensors_decoding_options.num_values_per_keypoint
)
tensors_decoding_options.xScale = self.tensors_decoding_options.x_scale
tensors_decoding_options.yScale = self.tensors_decoding_options.y_scale
tensors_decoding_options.wScale = self.tensors_decoding_options.w_scale
tensors_decoding_options.hScale = self.tensors_decoding_options.h_scale
tensors_decoding_options.applyExponentialOnBoxSize = (
self.tensors_decoding_options.apply_exponential_on_box_size
)
tensors_decoding_options.sigmoidScore = (
self.tensors_decoding_options.sigmoid_score
)
detector_options.tensorsDecodingOptions = tensors_decoding_options
# Get the object detector options flatbuffer.
b = flatbuffers.Builder(0)
b.Finish(detector_options.Pack(b), self._METADATA_FILE_IDENTIFIER)
detector_options_buf = b.Output()
# Add the object detector options flatbuffer in custom metadata.
custom_metadata = _metadata_fb.CustomMetadataT()
custom_metadata.name = self.name
custom_metadata.data = detector_options_buf
return custom_metadata
class MetadataWriter(metadata_writer.MetadataWriterBase):
"""MetadataWriter to write the metadata into the object detector."""
@classmethod
def create_for_models_with_nms(
cls,
model_buffer: bytearray,
input_norm_mean: List[float],
input_norm_std: List[float],
labels: metadata_writer.Labels,
score_calibration: Optional[metadata_writer.ScoreCalibration] = None,
) -> "MetadataWriter":
"""Creates MetadataWriter to write the metadata for object detector with postprocessing in the model.
This method create a metadata writer for the models with postprocessing [1].
The parameters required in this method are mandatory when using MediaPipe
Tasks.
Example usage:
metadata_writer = object_detector.Metadatawriter.create(model_buffer, ...)
tflite_content, json_content = metadata_writer.populate()
When calling `populate` function in this class, it returns TfLite content
and JSON content. Note that only the output TFLite is used for deployment.
The output JSON content is used to interpret the metadata content.
Args:
model_buffer: A valid flatbuffer loaded from the TFLite model file.
input_norm_mean: the mean value used in the input tensor normalization
[2].
input_norm_std: the std value used in the input tensor normalizarion [2].
labels: an instance of Labels helper class used in the output
classification tensor [3].
score_calibration: A container of the score calibration operation [4] in
the classification tensor. Optional if the model does not use score
calibration.
[1]:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/detection_postprocess.cc
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
[3]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
[4]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456
Returns:
A MetadataWriter object.
"""
writer = metadata_writer.MetadataWriter(model_buffer)
writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
writer.add_image_input(input_norm_mean, input_norm_std)
writer.add_detection_output(labels, score_calibration)
return cls(writer)
@classmethod
def create_for_models_without_nms(
cls,
model_buffer: bytearray,
input_norm_mean: List[float],
input_norm_std: List[float],
labels: metadata_writer.Labels,
ssd_anchors_options: SsdAnchorsOptions,
tensors_decoding_options: TensorsDecodingOptions,
output_tensors_order: metadata_info.RawDetectionOutputTensorsOrder = metadata_info.RawDetectionOutputTensorsOrder.UNSPECIFIED,
) -> "MetadataWriter":
"""Creates MetadataWriter to write the metadata for object detector without postprocessing in the model.
This method create a metadata writer for the models without postprocessing
[1].
The parameters required in this method are mandatory when using MediaPipe
Tasks.
Example usage:
metadata_writer = object_detector.Metadatawriter.create(model_buffer, ...)
tflite_content, json_content = metadata_writer.populate()
When calling `populate` function in this class, it returns TfLite content
and JSON content. Note that only the output TFLite is used for deployment.
The output JSON content is used to interpret the metadata content.
Args:
model_buffer: A valid flatbuffer loaded from the TFLite model file.
input_norm_mean: the mean value used in the input tensor normalization
[2].
input_norm_std: the std value used in the input tensor normalizarion [2].
labels: an instance of Labels helper class used in the output
classification tensor [3].
ssd_anchors_options: the ssd anchors options associated to the object
detector model.
tensors_decoding_options: the tensors decoding options used to decode the
object detector model output.
output_tensors_order: the order of the output tensors.
[1]:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/detection_postprocess.cc
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
[3]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
Returns:
A MetadataWriter object.
"""
writer = metadata_writer.MetadataWriter(model_buffer)
writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
writer.add_image_input(input_norm_mean, input_norm_std)
writer.add_raw_detection_output(
labels, output_tensors_order=output_tensors_order
)
option_md = ObjectDetectorOptionsMd(
ssd_anchors_options, tensors_decoding_options
)
writer.add_custom_metadata(option_md)
return cls(writer)
def populate(self) -> "tuple[bytearray, str]":
model_buf, _ = super().populate()
metadata_buf = metadata.get_metadata_buffer(model_buf)
json_content = convert_to_json(metadata_buf)
return model_buf, json_content
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/text_classifier.py
+119
View File
@@ -0,0 +1,119 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Writes metadata and label file to the Text classifier models."""
from typing import Union
from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
_MODEL_NAME = "TextClassifier"
_MODEL_DESCRIPTION = ("Classify the input text into a set of known categories.")
# The input tensor names of models created by Model Maker.
_DEFAULT_ID_NAME = "serving_default_input_word_ids:0"
_DEFAULT_MASK_NAME = "serving_default_input_mask:0"
_DEFAULT_SEGMENT_ID_NAME = "serving_default_input_type_ids:0"
class MetadataWriter(metadata_writer.MetadataWriterBase):
"""MetadataWriter to write the metadata into the text classifier."""
@classmethod
def create_for_regex_model(
cls, model_buffer: bytearray,
regex_tokenizer: metadata_writer.RegexTokenizer,
labels: metadata_writer.Labels) -> "MetadataWriter":
"""Creates MetadataWriter for TFLite model with regex tokentizer.
The parameters required in this method are mandatory when using MediaPipe
Tasks.
Note that only the output TFLite is used for deployment. The output JSON
content is used to interpret the metadata content.
Args:
model_buffer: A valid flatbuffer loaded from the TFLite model file.
regex_tokenizer: information of the regex tokenizer [1] used to process
the input string. If the tokenizer is `BertTokenizer` [2] or
`SentencePieceTokenizer` [3], please refer to
`create_for_bert_model`.
labels: an instance of Labels helper class used in the output
classification tensor [4].
[1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
[3]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
[4]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
Returns:
A MetadataWriter object.
"""
writer = metadata_writer.MetadataWriter(model_buffer)
writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
writer.add_regex_text_input(regex_tokenizer)
writer.add_classification_output(labels)
return cls(writer)
@classmethod
def create_for_bert_model(
cls,
model_buffer: bytearray,
tokenizer: Union[metadata_writer.BertTokenizer,
metadata_writer.SentencePieceTokenizer],
labels: metadata_writer.Labels,
ids_name: str = _DEFAULT_ID_NAME,
mask_name: str = _DEFAULT_MASK_NAME,
segment_name: str = _DEFAULT_SEGMENT_ID_NAME,
) -> "MetadataWriter":
"""Creates MetadataWriter for models with {Bert/SentencePiece}Tokenizer.
`ids_name`, `mask_name`, and `segment_name` correspond to the `Tensor.name`
in the TFLite schema, which help to determine the tensor order when
populating metadata. The default values come from Model Maker.
Args:
model_buffer: valid buffer of the model file.
tokenizer: information of the tokenizer used to process the input string,
if any. Supported tokenziers are: `BertTokenizer` [1] and
`SentencePieceTokenizer` [2]. If the tokenizer is `RegexTokenizer` [3],
refer to `create_for_regex_model`.
labels: an instance of Labels helper class used in the output
classification tensor [4].
ids_name: name of the ids tensor, which represents the tokenized ids of
the input text.
mask_name: name of the mask tensor, which represents the mask with `1` for
real tokens and `0` for padding tokens.
segment_name: name of the segment ids tensor, where `0` stands for the
first sequence, and `1` stands for the second sequence if exists. [1]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
[2]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
[3]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
[4]:
https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
Returns:
A MetadataWriter object.
"""
writer = metadata_writer.MetadataWriter(model_buffer)
writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
writer.add_bert_text_input(tokenizer, ids_name, mask_name, segment_name)
writer.add_classification_output(labels)
return cls(writer)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py
+91
View File
@@ -0,0 +1,91 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
# ==============================================================================
"""Helper methods for writing metadata into TFLite models."""
from typing import List
from mediapipe.tasks.metadata import schema_py_generated as _schema_fb
def get_input_tensor_names(model_buffer: bytearray) -> List[str]:
"""Gets a list of the input tensor names."""
subgraph = get_subgraph(model_buffer)
tensor_names = []
for i in range(subgraph.InputsLength()):
index = subgraph.Inputs(i)
tensor_names.append(subgraph.Tensors(index).Name().decode("utf-8"))
return tensor_names
def get_output_tensor_names(model_buffer: bytearray) -> List[str]:
"""Gets a list of the output tensor names."""
subgraph = get_subgraph(model_buffer)
tensor_names = []
for i in range(subgraph.OutputsLength()):
index = subgraph.Outputs(i)
tensor_names.append(subgraph.Tensors(index).Name().decode("utf-8"))
return tensor_names
def get_input_tensor_types(
model_buffer: bytearray) -> List[_schema_fb.TensorType]:
"""Gets a list of the input tensor types."""
subgraph = get_subgraph(model_buffer)
tensor_types = []
for i in range(subgraph.InputsLength()):
index = subgraph.Inputs(i)
tensor_types.append(subgraph.Tensors(index).Type())
return tensor_types
def get_output_tensor_types(
model_buffer: bytearray) -> List[_schema_fb.TensorType]:
"""Gets a list of the output tensor types."""
subgraph = get_subgraph(model_buffer)
tensor_types = []
for i in range(subgraph.OutputsLength()):
index = subgraph.Outputs(i)
tensor_types.append(subgraph.Tensors(index).Type())
return tensor_types
def get_output_tensor_indices(model_buffer: bytearray) -> List[int]:
"""Gets a list of the output tensor indices."""
subgraph = get_subgraph(model_buffer)
return subgraph.OutputsAsNumpy()
def get_subgraph(model_buffer: bytearray) -> _schema_fb.SubGraph:
"""Gets the subgraph of the model.
TFLite does not support multi-subgraph. A model should have exactly one
subgraph.
Args:
model_buffer: valid buffer of the model file.
Returns:
The subgraph of the model.
Raises:
ValueError: if the model has more than one subgraph or has no subgraph.
"""
model = _schema_fb.Model.GetRootAsModel(model_buffer, 0)
# Use the first subgraph as default. TFLite Interpreter doesn't support
# multiple subgraphs yet, but models with mini-benchmark may have multiple
# subgraphs for acceleration evaluation purpose.
return model.Subgraphs(0)
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/audio/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/audio/audio_classifier_test.py
+405
View File
@@ -0,0 +1,405 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Tests for audio classifier."""
import os
from typing import List, Tuple
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from scipy.io import wavfile
from mediapipe.tasks.python.audio import audio_classifier
from mediapipe.tasks.python.audio.core import audio_record
from mediapipe.tasks.python.audio.core import audio_task_running_mode
from mediapipe.tasks.python.components.containers import audio_data as audio_data_module
from mediapipe.tasks.python.components.containers import classification_result as classification_result_module
from mediapipe.tasks.python.core import base_options as base_options_module
from mediapipe.tasks.python.test import test_utils
_AudioClassifier = audio_classifier.AudioClassifier
_AudioClassifierOptions = audio_classifier.AudioClassifierOptions
_AudioClassifierResult = classification_result_module.ClassificationResult
_AudioData = audio_data_module.AudioData
_AudioRecord = audio_record.AudioRecord
_BaseOptions = base_options_module.BaseOptions
_RUNNING_MODE = audio_task_running_mode.AudioTaskRunningMode
_YAMNET_MODEL_FILE = 'yamnet_audio_classifier_with_metadata.tflite'
_YAMNET_MODEL_SAMPLE_RATE = 16000
_TWO_HEADS_MODEL_FILE = 'two_heads.tflite'
_SPEECH_WAV_16K_MONO = 'speech_16000_hz_mono.wav'
_SPEECH_WAV_48K_MONO = 'speech_48000_hz_mono.wav'
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/audio'
_TWO_HEADS_WAV_16K_MONO = 'two_heads_16000_hz_mono.wav'
_TWO_HEADS_WAV_44K_MONO = 'two_heads_44100_hz_mono.wav'
_YAMNET_NUM_OF_SAMPLES = 15600
_MILLISECONDS_PER_SECOND = 1000
class AudioClassifierTest(parameterized.TestCase):
def setUp(self):
super().setUp()
self.yamnet_model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, _YAMNET_MODEL_FILE))
self.two_heads_model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, _TWO_HEADS_MODEL_FILE))
def _read_wav_file(self, file_name) -> _AudioData:
sample_rate, buffer = wavfile.read(
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, file_name)))
return _AudioData.create_from_array(
buffer.astype(float) / np.iinfo(np.int16).max, sample_rate)
def _read_wav_file_as_stream(self, file_name) -> List[Tuple[_AudioData, int]]:
sample_rate, buffer = wavfile.read(
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, file_name)))
audio_data_list = []
start = 0
step_size = _YAMNET_NUM_OF_SAMPLES * sample_rate / _YAMNET_MODEL_SAMPLE_RATE
while start < len(buffer):
end = min(start + (int)(step_size), len(buffer))
audio_data_list.append((_AudioData.create_from_array(
buffer[start:end].astype(float) / np.iinfo(np.int16).max,
sample_rate), (int)(start / sample_rate * _MILLISECONDS_PER_SECOND)))
start = end
return audio_data_list
# TODO: Compares the exact score values to capture unexpected
# changes in the inference pipeline.
def _check_yamnet_result(
self,
classification_result_list: List[_AudioClassifierResult],
expected_num_categories=521):
self.assertLen(classification_result_list, 5)
for idx, timestamp in enumerate([0, 975, 1950, 2925]):
classification_result = classification_result_list[idx]
self.assertEqual(classification_result.timestamp_ms, timestamp)
self.assertLen(classification_result.classifications, 1)
classifcation = classification_result.classifications[0]
self.assertEqual(classifcation.head_index, 0)
self.assertEqual(classifcation.head_name, 'scores')
self.assertLen(classifcation.categories, expected_num_categories)
audio_category = classifcation.categories[0]
self.assertEqual(audio_category.index, 0)
self.assertEqual(audio_category.category_name, 'Speech')
self.assertGreater(audio_category.score, 0.9)
# TODO: Compares the exact score values to capture unexpected
# changes in the inference pipeline.
def _check_two_heads_result(
self,
classification_result_list: List[_AudioClassifierResult],
first_head_expected_num_categories=521,
second_head_expected_num_categories=5):
self.assertGreaterEqual(len(classification_result_list), 1)
self.assertLessEqual(len(classification_result_list), 2)
# Checks the first result.
classification_result = classification_result_list[0]
self.assertEqual(classification_result.timestamp_ms, 0)
self.assertLen(classification_result.classifications, 2)
# Checks the first head.
yamnet_classifcation = classification_result.classifications[0]
self.assertEqual(yamnet_classifcation.head_index, 0)
self.assertEqual(yamnet_classifcation.head_name, 'yamnet_classification')
self.assertLen(yamnet_classifcation.categories,
first_head_expected_num_categories)
# Checks the second head.
yamnet_category = yamnet_classifcation.categories[0]
self.assertEqual(yamnet_category.index, 508)
self.assertEqual(yamnet_category.category_name, 'Environmental noise')
self.assertGreater(yamnet_category.score, 0.5)
bird_classifcation = classification_result.classifications[1]
self.assertEqual(bird_classifcation.head_index, 1)
self.assertEqual(bird_classifcation.head_name, 'bird_classification')
self.assertLen(bird_classifcation.categories,
second_head_expected_num_categories)
bird_category = bird_classifcation.categories[0]
self.assertEqual(bird_category.index, 4)
self.assertEqual(bird_category.category_name, 'Chestnut-crowned Antpitta')
self.assertGreater(bird_category.score, 0.93)
# Checks the second result, if present.
if len(classification_result_list) == 2:
classification_result = classification_result_list[1]
self.assertEqual(classification_result.timestamp_ms, 975)
self.assertLen(classification_result.classifications, 2)
# Checks the first head.
yamnet_classifcation = classification_result.classifications[0]
self.assertEqual(yamnet_classifcation.head_index, 0)
self.assertEqual(yamnet_classifcation.head_name, 'yamnet_classification')
self.assertLen(yamnet_classifcation.categories,
first_head_expected_num_categories)
yamnet_category = yamnet_classifcation.categories[0]
self.assertEqual(yamnet_category.index, 494)
self.assertEqual(yamnet_category.category_name, 'Silence')
self.assertGreater(yamnet_category.score, 0.9)
bird_classifcation = classification_result.classifications[1]
self.assertEqual(bird_classifcation.head_index, 1)
self.assertEqual(bird_classifcation.head_name, 'bird_classification')
self.assertLen(bird_classifcation.categories,
second_head_expected_num_categories)
# Checks the second head.
bird_category = bird_classifcation.categories[0]
self.assertEqual(bird_category.index, 1)
self.assertEqual(bird_category.category_name, 'White-breasted Wood-Wren')
self.assertGreater(bird_category.score, 0.99)
def test_create_from_file_succeeds_with_valid_model_path(self):
# Creates with default option and valid model file successfully.
with _AudioClassifier.create_from_model_path(
self.yamnet_model_path) as classifier:
self.assertIsInstance(classifier, _AudioClassifier)
def test_create_from_options_succeeds_with_valid_model_path(self):
# Creates with options containing model file successfully.
with _AudioClassifier.create_from_options(
_AudioClassifierOptions(
base_options=_BaseOptions(
model_asset_path=self.yamnet_model_path))) as classifier:
self.assertIsInstance(classifier, _AudioClassifier)
def test_create_from_options_fails_with_invalid_model_path(self):
with self.assertRaisesRegex(
FileNotFoundError,
'Unable to open file at /path/to/invalid/model.tflite',
):
base_options = _BaseOptions(
model_asset_path='/path/to/invalid/model.tflite')
options = _AudioClassifierOptions(base_options=base_options)
classifier = _AudioClassifier.create_from_options(options)
classifier.close()
def test_create_from_options_succeeds_with_valid_model_content(self):
# Creates with options containing model content successfully.
with open(self.yamnet_model_path, 'rb') as f:
base_options = _BaseOptions(model_asset_buffer=f.read())
options = _AudioClassifierOptions(base_options=base_options)
classifier = _AudioClassifier.create_from_options(options)
self.assertIsInstance(classifier, _AudioClassifier)
classifier.close()
@parameterized.parameters((_SPEECH_WAV_16K_MONO), (_SPEECH_WAV_48K_MONO))
def test_classify_with_yamnet_model(self, audio_file):
with _AudioClassifier.create_from_model_path(
self.yamnet_model_path) as classifier:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_yamnet_result(classification_result_list)
def test_classify_with_yamnet_model_and_inputs_at_different_sample_rates(
self):
with _AudioClassifier.create_from_model_path(
self.yamnet_model_path) as classifier:
for audio_file in [_SPEECH_WAV_16K_MONO, _SPEECH_WAV_48K_MONO]:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_yamnet_result(classification_result_list)
@mock.patch('sounddevice.InputStream', return_value=mock.MagicMock())
def test_create_audio_record_from_classifier_succeeds(self, _):
# Creates AudioRecord instance using the classifier successfully.
with _AudioClassifier.create_from_model_path(
self.yamnet_model_path
) as classifier:
self.assertIsInstance(classifier, _AudioClassifier)
record = classifier.create_audio_record(1, 16000, 16000)
self.assertIsInstance(record, _AudioRecord)
self.assertEqual(record.channels, 1)
self.assertEqual(record.sampling_rate, 16000)
self.assertEqual(record.buffer_size, 16000)
def test_max_result_options(self):
with _AudioClassifier.create_from_options(
_AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
max_results=1)) as classifier:
for audio_file in [_SPEECH_WAV_16K_MONO, _SPEECH_WAV_16K_MONO]:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_yamnet_result(
classification_result_list, expected_num_categories=1)
def test_score_threshold_options(self):
with _AudioClassifier.create_from_options(
_AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
score_threshold=0.9)) as classifier:
for audio_file in [_SPEECH_WAV_16K_MONO, _SPEECH_WAV_16K_MONO]:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_yamnet_result(
classification_result_list, expected_num_categories=1)
def test_allow_list_option(self):
with _AudioClassifier.create_from_options(
_AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
category_allowlist=['Speech'])) as classifier:
for audio_file in [_SPEECH_WAV_16K_MONO, _SPEECH_WAV_16K_MONO]:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_yamnet_result(
classification_result_list, expected_num_categories=1)
def test_combined_allowlist_and_denylist(self):
# Fails with combined allowlist and denylist
with self.assertRaisesRegex(
ValueError,
'`category_allowlist` and `category_denylist` are mutually exclusive'
' options',
):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
category_allowlist=['foo'],
category_denylist=['bar'])
with _AudioClassifier.create_from_options(options) as unused_classifier:
pass
@parameterized.parameters((_TWO_HEADS_WAV_16K_MONO),
(_TWO_HEADS_WAV_44K_MONO))
def test_classify_with_two_heads_model_and_inputs_at_different_sample_rates(
self, audio_file):
with _AudioClassifier.create_from_model_path(
self.two_heads_model_path) as classifier:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_two_heads_result(classification_result_list)
def test_classify_with_two_heads_model(self):
with _AudioClassifier.create_from_model_path(
self.two_heads_model_path) as classifier:
for audio_file in [_TWO_HEADS_WAV_16K_MONO, _TWO_HEADS_WAV_44K_MONO]:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_two_heads_result(classification_result_list)
def test_classify_with_two_heads_model_with_max_results(self):
with _AudioClassifier.create_from_options(
_AudioClassifierOptions(
base_options=_BaseOptions(
model_asset_path=self.two_heads_model_path),
max_results=1)) as classifier:
for audio_file in [_TWO_HEADS_WAV_16K_MONO, _TWO_HEADS_WAV_44K_MONO]:
classification_result_list = classifier.classify(
self._read_wav_file(audio_file))
self._check_two_heads_result(classification_result_list, 1, 1)
def test_missing_sample_rate_in_audio_clips_mode(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_CLIPS)
with self.assertRaisesRegex(ValueError,
r'Must provide the audio sample rate'):
with _AudioClassifier.create_from_options(options) as classifier:
classifier.classify(_AudioData(buffer_length=100))
def test_missing_sample_rate_in_audio_stream_mode(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_STREAM,
result_callback=mock.MagicMock())
with self.assertRaisesRegex(ValueError,
r'provide the audio sample rate in audio data'):
with _AudioClassifier.create_from_options(options) as classifier:
classifier.classify(_AudioData(buffer_length=100))
def test_missing_result_callback(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_STREAM)
with self.assertRaisesRegex(
ValueError,
'The audio task is in audio stream mode, a user-defined result callback'
' must be provided.',
):
with _AudioClassifier.create_from_options(options) as unused_classifier:
pass
def test_illegal_result_callback(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_CLIPS,
result_callback=mock.MagicMock())
with self.assertRaisesRegex(
ValueError,
'The audio task is in audio clips mode, a user-defined result callback'
" shouldn't be provided",
):
with _AudioClassifier.create_from_options(options) as unused_classifier:
pass
def test_calling_classify_in_audio_stream_mode(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_STREAM,
result_callback=mock.MagicMock())
with _AudioClassifier.create_from_options(options) as classifier:
with self.assertRaisesRegex(
ValueError,
'Task is not initialized with the audio clips mode. Current running'
' mode:audio stream mode',
):
classifier.classify(self._read_wav_file(_SPEECH_WAV_16K_MONO))
def test_calling_classify_async_in_audio_clips_mode(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_CLIPS)
with _AudioClassifier.create_from_options(options) as classifier:
with self.assertRaisesRegex(
ValueError,
'Task is not initialized with the audio stream mode. Current running'
' mode:audio clips mode',
):
classifier.classify_async(self._read_wav_file(_SPEECH_WAV_16K_MONO), 0)
def test_classify_async_calls_with_illegal_timestamp(self):
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_STREAM,
result_callback=mock.MagicMock())
with _AudioClassifier.create_from_options(options) as classifier:
classifier.classify_async(self._read_wav_file(_SPEECH_WAV_16K_MONO), 100)
with self.assertRaisesRegex(
ValueError, 'Input timestamp must be monotonically increasing.'
):
classifier.classify_async(self._read_wav_file(_SPEECH_WAV_16K_MONO), 0)
@parameterized.parameters((_SPEECH_WAV_16K_MONO), (_SPEECH_WAV_48K_MONO))
def test_classify_async(self, audio_file):
classification_result_list = []
def save_result(result: _AudioClassifierResult, timestamp_ms: int):
result.timestamp_ms = timestamp_ms
classification_result_list.append(result)
options = _AudioClassifierOptions(
base_options=_BaseOptions(model_asset_path=self.yamnet_model_path),
running_mode=_RUNNING_MODE.AUDIO_STREAM,
max_results=1,
result_callback=save_result)
classifier = _AudioClassifier.create_from_options(options)
audio_data_list = self._read_wav_file_as_stream(audio_file)
for audio_data, timestamp_ms in audio_data_list:
classifier.classify_async(audio_data, timestamp_ms)
classifier.close()
self._check_yamnet_result(
classification_result_list, expected_num_categories=1)
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/test_utils.py
+196
View File
@@ -0,0 +1,196 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Test util for MediaPipe Tasks."""
import difflib
import os
from absl import flags
import six
from google.protobuf import descriptor
from google.protobuf import descriptor_pool
from google.protobuf import text_format
from mediapipe.python._framework_bindings import image as image_module
from mediapipe.python._framework_bindings import image_frame as image_frame_module
FLAGS = flags.FLAGS
_Image = image_module.Image
_ImageFormat = image_frame_module.ImageFormat
_RGB_CHANNELS = 3
def test_srcdir():
"""Returns the path where to look for test data files."""
if "test_srcdir" in flags.FLAGS:
return flags.FLAGS["test_srcdir"].value
elif "TEST_SRCDIR" in os.environ:
return os.environ["TEST_SRCDIR"]
else:
raise RuntimeError("Missing TEST_SRCDIR environment.")
def get_test_data_path(file_or_dirname_path: str) -> str:
"""Returns full test data path."""
for directory, subdirs, files in os.walk(test_srcdir()):
for f in subdirs + files:
path = os.path.join(directory, f)
if path.endswith(file_or_dirname_path):
return path
raise ValueError(
"No %s in test directory: %s." % (file_or_dirname_path, test_srcdir())
)
def create_calibration_file(
file_dir: str,
file_name: str = "score_calibration.txt",
content: str = "1.0,2.0,3.0,4.0",
) -> str:
"""Creates the calibration file."""
calibration_file = os.path.join(file_dir, file_name)
with open(calibration_file, mode="w") as file:
file.write(content)
return calibration_file
def assert_proto_equals(
self, a, b, check_initialized=True, normalize_numbers=True, msg=None
):
"""assert_proto_equals() is useful for unit tests.
It produces much more helpful output than assertEqual() for proto2 messages.
Fails with a useful error if a and b aren't equal. Comparison of repeated
fields matches the semantics of unittest.TestCase.assertEqual(), ie order and
extra duplicates fields matter.
This is a fork of https://github.com/tensorflow/tensorflow/blob/
master/tensorflow/python/util/protobuf/compare.py#L73. We use slightly
different rounding cutoffs to support Mac usage.
Args:
self: absltest.testing.parameterized.TestCase
a: proto2 PB instance, or text string representing one.
b: proto2 PB instance -- message.Message or subclass thereof.
check_initialized: boolean, whether to fail if either a or b isn't
initialized.
normalize_numbers: boolean, whether to normalize types and precision of
numbers before comparison.
msg: if specified, is used as the error message on failure.
"""
pool = descriptor_pool.Default()
if isinstance(a, six.string_types):
a = text_format.Parse(a, b.__class__(), descriptor_pool=pool)
for pb in a, b:
if check_initialized:
errors = pb.FindInitializationErrors()
if errors:
self.fail("Initialization errors: %s\n%s" % (errors, pb))
if normalize_numbers:
_normalize_number_fields(pb)
a_str = text_format.MessageToString(a, descriptor_pool=pool)
b_str = text_format.MessageToString(b, descriptor_pool=pool)
# Some Python versions would perform regular diff instead of multi-line
# diff if string is longer than 2**16. We substitute this behavior
# with a call to unified_diff instead to have easier-to-read diffs.
# For context, see: https://bugs.python.org/issue11763.
if len(a_str) < 2**16 and len(b_str) < 2**16:
self.assertMultiLineEqual(a_str, b_str, msg=msg)
else:
diff = "".join(
difflib.unified_diff(a_str.splitlines(True), b_str.splitlines(True))
)
if diff:
self.fail("%s :\n%s" % (msg, diff))
def _normalize_number_fields(pb):
"""Normalizes types and precisions of number fields in a protocol buffer.
Due to subtleties in the python protocol buffer implementation, it is possible
for values to have different types and precision depending on whether they
were set and retrieved directly or deserialized from a protobuf. This function
normalizes integer values to ints and longs based on width, 32-bit floats to
five digits of precision to account for python always storing them as 64-bit,
and ensures doubles are floating point for when they're set to integers.
Modifies pb in place. Recurses into nested objects. https://github.com/tensorf
low/tensorflow/blob/master/tensorflow/python/util/protobuf/compare.py#L118
Args:
pb: proto2 message.
Returns:
the given pb, modified in place.
"""
for desc, values in pb.ListFields():
is_repeated = True
if desc.label != descriptor.FieldDescriptor.LABEL_REPEATED:
is_repeated = False
values = [values]
normalized_values = None
# We force 32-bit values to int and 64-bit values to long to make
# alternate implementations where the distinction is more significant
# (e.g. the C++ implementation) simpler.
if desc.type in (
descriptor.FieldDescriptor.TYPE_INT64,
descriptor.FieldDescriptor.TYPE_UINT64,
descriptor.FieldDescriptor.TYPE_SINT64,
):
normalized_values = [int(x) for x in values]
elif desc.type in (
descriptor.FieldDescriptor.TYPE_INT32,
descriptor.FieldDescriptor.TYPE_UINT32,
descriptor.FieldDescriptor.TYPE_SINT32,
descriptor.FieldDescriptor.TYPE_ENUM,
):
normalized_values = [int(x) for x in values]
elif desc.type == descriptor.FieldDescriptor.TYPE_FLOAT:
normalized_values = [round(x, 4) for x in values]
elif desc.type == descriptor.FieldDescriptor.TYPE_DOUBLE:
normalized_values = [round(float(x), 6) for x in values]
if normalized_values is not None:
if is_repeated:
pb.ClearField(desc.name)
getattr(pb, desc.name).extend(normalized_values)
else:
setattr(pb, desc.name, normalized_values[0])
if (
desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE
or desc.type == descriptor.FieldDescriptor.TYPE_GROUP
):
if (
desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE
and desc.message_type.has_options
and desc.message_type.GetOptions().map_entry
):
# This is a map, only recurse if the values have a message type.
if (
desc.message_type.fields_by_number[2].type
== descriptor.FieldDescriptor.TYPE_MESSAGE
):
for v in six.itervalues(values):
_normalize_number_fields(v)
else:
for v in values:
# recursive step
_normalize_number_fields(v)
return pb
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/text/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/text/language_detector_test.py
+226
View File
@@ -0,0 +1,226 @@
# Copyright 2023 The MediaPipe Authors.
#
# 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.
"""Tests for language detector."""
import enum
import os
from absl.testing import absltest
from absl.testing import parameterized
from mediapipe.tasks.python.components.containers import category
from mediapipe.tasks.python.components.containers import classification_result as classification_result_module
from mediapipe.tasks.python.core import base_options as base_options_module
from mediapipe.tasks.python.test import test_utils
from mediapipe.tasks.python.text import language_detector
LanguageDetectorResult = language_detector.LanguageDetectorResult
LanguageDetectorPrediction = (
language_detector.LanguageDetectorResult.Detection
)
_BaseOptions = base_options_module.BaseOptions
_Category = category.Category
_Classifications = classification_result_module.Classifications
_LanguageDetector = language_detector.LanguageDetector
_LanguageDetectorOptions = language_detector.LanguageDetectorOptions
_LANGUAGE_DETECTOR_MODEL = "language_detector.tflite"
_TEST_DATA_DIR = "mediapipe/tasks/testdata/text"
_SCORE_THRESHOLD = 0.3
_EN_TEXT = "To be, or not to be, that is the question"
_EN_EXPECTED_RESULT = LanguageDetectorResult(
[LanguageDetectorPrediction("en", 0.999856)]
)
_FR_TEXT = (
"Il y a beaucoup de bouches qui parlent et fort peu de têtes qui pensent."
)
_FR_EXPECTED_RESULT = LanguageDetectorResult(
[LanguageDetectorPrediction("fr", 0.999781)]
)
_RU_TEXT = "это какой-то английский язык"
_RU_EXPECTED_RESULT = LanguageDetectorResult(
[LanguageDetectorPrediction("ru", 0.993362)]
)
_MIXED_TEXT = "分久必合合久必分"
_MIXED_EXPECTED_RESULT = LanguageDetectorResult([
LanguageDetectorPrediction("zh", 0.505424),
LanguageDetectorPrediction("ja", 0.481617),
])
_TOLERANCE = 1e-6
class ModelFileType(enum.Enum):
FILE_CONTENT = 1
FILE_NAME = 2
class LanguageDetectorTest(parameterized.TestCase):
def setUp(self):
super().setUp()
self.model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, _LANGUAGE_DETECTOR_MODEL)
)
def _expect_language_detector_result_correct(
self,
actual_result: LanguageDetectorResult,
expect_result: LanguageDetectorResult,
):
for i, prediction in enumerate(actual_result.detections):
expected_prediction = expect_result.detections[i]
self.assertEqual(
prediction.language_code,
expected_prediction.language_code,
)
self.assertAlmostEqual(
prediction.probability,
expected_prediction.probability,
delta=_TOLERANCE,
)
def test_create_from_file_succeeds_with_valid_model_path(self):
# Creates with default option and valid model file successfully.
with _LanguageDetector.create_from_model_path(self.model_path) as detector:
self.assertIsInstance(detector, _LanguageDetector)
def test_create_from_options_succeeds_with_valid_model_path(self):
# Creates with options containing model file successfully.
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _LanguageDetectorOptions(base_options=base_options)
with _LanguageDetector.create_from_options(options) as detector:
self.assertIsInstance(detector, _LanguageDetector)
def test_create_from_options_fails_with_invalid_model_path(self):
with self.assertRaises(FileNotFoundError):
base_options = _BaseOptions(
model_asset_path="/path/to/invalid/model.tflite"
)
options = _LanguageDetectorOptions(base_options=base_options)
_LanguageDetector.create_from_options(options)
def test_create_from_options_succeeds_with_valid_model_content(self):
# Creates with options containing model content successfully.
with open(self.model_path, "rb") as f:
base_options = _BaseOptions(model_asset_buffer=f.read())
options = _LanguageDetectorOptions(base_options=base_options)
detector = _LanguageDetector.create_from_options(options)
self.assertIsInstance(detector, _LanguageDetector)
@parameterized.parameters(
(ModelFileType.FILE_NAME, _EN_TEXT, _EN_EXPECTED_RESULT),
(ModelFileType.FILE_CONTENT, _EN_TEXT, _EN_EXPECTED_RESULT),
(ModelFileType.FILE_NAME, _FR_TEXT, _FR_EXPECTED_RESULT),
(ModelFileType.FILE_CONTENT, _FR_TEXT, _FR_EXPECTED_RESULT),
(ModelFileType.FILE_NAME, _RU_TEXT, _RU_EXPECTED_RESULT),
(ModelFileType.FILE_CONTENT, _RU_TEXT, _RU_EXPECTED_RESULT),
(ModelFileType.FILE_NAME, _MIXED_TEXT, _MIXED_EXPECTED_RESULT),
(ModelFileType.FILE_CONTENT, _MIXED_TEXT, _MIXED_EXPECTED_RESULT),
)
def test_detect(self, model_file_type, text, expected_result):
# Creates detector.
if model_file_type is ModelFileType.FILE_NAME:
base_options = _BaseOptions(model_asset_path=self.model_path)
elif model_file_type is ModelFileType.FILE_CONTENT:
with open(self.model_path, "rb") as f:
model_content = f.read()
base_options = _BaseOptions(model_asset_buffer=model_content)
else:
# Should never happen
raise ValueError("model_file_type is invalid.")
options = _LanguageDetectorOptions(
base_options=base_options, score_threshold=_SCORE_THRESHOLD
)
detector = _LanguageDetector.create_from_options(options)
# Performs language detection on the input.
text_result = detector.detect(text)
# Comparing results.
self._expect_language_detector_result_correct(text_result, expected_result)
# Closes the detector explicitly when the detector is not used in
# a context.
detector.close()
@parameterized.parameters(
(ModelFileType.FILE_NAME, _EN_TEXT, _EN_EXPECTED_RESULT),
(ModelFileType.FILE_NAME, _FR_TEXT, _FR_EXPECTED_RESULT),
(ModelFileType.FILE_NAME, _RU_TEXT, _RU_EXPECTED_RESULT),
(ModelFileType.FILE_CONTENT, _MIXED_TEXT, _MIXED_EXPECTED_RESULT),
)
def test_detect_in_context(self, model_file_type, text, expected_result):
# Creates detector.
if model_file_type is ModelFileType.FILE_NAME:
base_options = _BaseOptions(model_asset_path=self.model_path)
elif model_file_type is ModelFileType.FILE_CONTENT:
with open(self.model_path, "rb") as f:
model_content = f.read()
base_options = _BaseOptions(model_asset_buffer=model_content)
else:
# Should never happen
raise ValueError("model_file_type is invalid.")
options = _LanguageDetectorOptions(
base_options=base_options, score_threshold=_SCORE_THRESHOLD
)
with _LanguageDetector.create_from_options(options) as detector:
# Performs language detection on the input.
text_result = detector.detect(text)
# Comparing results.
self._expect_language_detector_result_correct(
text_result, expected_result
)
def test_allowlist_option(self):
# Creates detector.
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _LanguageDetectorOptions(
base_options=base_options,
score_threshold=_SCORE_THRESHOLD,
category_allowlist=["ja"],
)
with _LanguageDetector.create_from_options(options) as detector:
# Performs language detection on the input.
text_result = detector.detect(_MIXED_TEXT)
# Comparing results.
expected_result = LanguageDetectorResult(
[LanguageDetectorPrediction("ja", 0.481617)]
)
self._expect_language_detector_result_correct(
text_result, expected_result
)
def test_denylist_option(self):
# Creates detector.
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _LanguageDetectorOptions(
base_options=base_options,
score_threshold=_SCORE_THRESHOLD,
category_denylist=["ja"],
)
with _LanguageDetector.create_from_options(options) as detector:
# Performs language detection on the input.
text_result = detector.detect(_MIXED_TEXT)
# Comparing results.
expected_result = LanguageDetectorResult(
[LanguageDetectorPrediction("zh", 0.505424)]
)
self._expect_language_detector_result_correct(
text_result, expected_result
)
if __name__ == "__main__":
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/text/text_classifier_test.py
+295
View File
@@ -0,0 +1,295 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Tests for text classifier."""
import enum
import os
from absl.testing import absltest
from absl.testing import parameterized
from mediapipe.tasks.python.components.containers import category
from mediapipe.tasks.python.components.containers import classification_result as classification_result_module
from mediapipe.tasks.python.core import base_options as base_options_module
from mediapipe.tasks.python.test import test_utils
from mediapipe.tasks.python.text import text_classifier
TextClassifierResult = classification_result_module.ClassificationResult
_BaseOptions = base_options_module.BaseOptions
_Category = category.Category
_Classifications = classification_result_module.Classifications
_TextClassifier = text_classifier.TextClassifier
_TextClassifierOptions = text_classifier.TextClassifierOptions
_BERT_MODEL_FILE = 'bert_text_classifier.tflite'
_REGEX_MODEL_FILE = 'test_model_text_classifier_with_regex_tokenizer.tflite'
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
_NEGATIVE_TEXT = 'What a waste of my time.'
_POSITIVE_TEXT = ('This is the best movie I’ve seen in recent years.'
'Strongly recommend it!')
_BERT_NEGATIVE_RESULTS = TextClassifierResult(
classifications=[
_Classifications(
categories=[
_Category(
index=0,
score=0.9995,
display_name='',
category_name='negative'),
_Category(
index=1,
score=0.0005,
display_name='',
category_name='positive')
],
head_index=0,
head_name='probability')
],
timestamp_ms=0)
_BERT_POSITIVE_RESULTS = TextClassifierResult(
classifications=[
_Classifications(
categories=[
_Category(
index=1,
score=0.9994,
display_name='',
category_name='positive',
),
_Category(
index=0,
score=0.0006,
display_name='',
category_name='negative',
),
],
head_index=0,
head_name='probability',
)
],
timestamp_ms=0,
)
_REGEX_NEGATIVE_RESULTS = TextClassifierResult(
classifications=[
_Classifications(
categories=[
_Category(
index=0,
score=0.81313,
display_name='',
category_name='Negative'),
_Category(
index=1,
score=0.1868704,
display_name='',
category_name='Positive')
],
head_index=0,
head_name='probability')
],
timestamp_ms=0)
_REGEX_POSITIVE_RESULTS = TextClassifierResult(
classifications=[
_Classifications(
categories=[
_Category(
index=1,
score=0.5134273,
display_name='',
category_name='Positive'),
_Category(
index=0,
score=0.486573,
display_name='',
category_name='Negative')
],
head_index=0,
head_name='probability')
],
timestamp_ms=0)
class ModelFileType(enum.Enum):
FILE_CONTENT = 1
FILE_NAME = 2
class TextClassifierTest(parameterized.TestCase):
def setUp(self):
super().setUp()
self.model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, _BERT_MODEL_FILE))
def assertTextClassifierResultEquals(
self,
result: TextClassifierResult,
expected_result: TextClassifierResult
):
self.assertEqual(result.timestamp_ms, expected_result.timestamp_ms)
self.assertLen(result.classifications, len(expected_result.classifications))
for i, actual_classification in enumerate(result.classifications):
expected_classification = expected_result.classifications[i]
self.assertEqual(
actual_classification.head_index, expected_classification.head_index
)
self.assertEqual(
actual_classification.head_name, expected_classification.head_name
)
self.assertLen(
actual_classification.categories,
len(expected_classification.categories),
)
for j, actual_category in enumerate(actual_classification.categories):
expected_category = expected_classification.categories[j]
self.assertEqual(actual_category.index, expected_category.index)
self.assertEqual(
actual_category.display_name or '', expected_category.display_name
)
self.assertEqual(
actual_category.category_name, expected_category.category_name
)
self.assertAlmostEqual(
actual_category.score, expected_category.score, delta=1e-4
)
def test_create_from_file_succeeds_with_valid_model_path(self):
# Creates with default option and valid model file successfully.
with _TextClassifier.create_from_model_path(self.model_path) as classifier:
self.assertIsInstance(classifier, _TextClassifier)
def test_create_from_options_succeeds_with_valid_model_path(self):
# Creates with options containing model file successfully.
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _TextClassifierOptions(base_options=base_options)
with _TextClassifier.create_from_options(options) as classifier:
self.assertIsInstance(classifier, _TextClassifier)
def test_create_from_options_fails_with_invalid_model_path(self):
with self.assertRaisesRegex(
FileNotFoundError,
'Unable to open file at /path/to/invalid/model.tflite',
):
base_options = _BaseOptions(
model_asset_path='/path/to/invalid/model.tflite')
options = _TextClassifierOptions(base_options=base_options)
_TextClassifier.create_from_options(options)
def test_create_from_options_succeeds_with_valid_model_content(self):
# Creates with options containing model content successfully.
with open(self.model_path, 'rb') as f:
base_options = _BaseOptions(model_asset_buffer=f.read())
options = _TextClassifierOptions(base_options=base_options)
classifier = _TextClassifier.create_from_options(options)
self.assertIsInstance(classifier, _TextClassifier)
def test_create_from_options_succeeds_with_allow_list(self):
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _TextClassifierOptions(
base_options=base_options, category_allowlist=['positive']
)
with _TextClassifier.create_from_options(options) as classifier:
self.assertIsInstance(classifier, _TextClassifier)
def test_create_from_options_succeeds_with_deny_list(self):
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _TextClassifierOptions(
base_options=base_options, category_denylist=['negative']
)
with _TextClassifier.create_from_options(options) as classifier:
self.assertIsInstance(classifier, _TextClassifier)
def test_create_from_options_succeeds_with_display_names_locale(self):
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _TextClassifierOptions(
base_options=base_options, display_names_locale='en'
)
with _TextClassifier.create_from_options(options) as classifier:
self.assertIsInstance(classifier, _TextClassifier)
@parameterized.parameters(
(ModelFileType.FILE_NAME, _BERT_MODEL_FILE, _NEGATIVE_TEXT,
_BERT_NEGATIVE_RESULTS), (ModelFileType.FILE_CONTENT, _BERT_MODEL_FILE,
_NEGATIVE_TEXT, _BERT_NEGATIVE_RESULTS),
(ModelFileType.FILE_NAME, _BERT_MODEL_FILE, _POSITIVE_TEXT,
_BERT_POSITIVE_RESULTS), (ModelFileType.FILE_CONTENT, _BERT_MODEL_FILE,
_POSITIVE_TEXT, _BERT_POSITIVE_RESULTS),
(ModelFileType.FILE_NAME, _REGEX_MODEL_FILE, _NEGATIVE_TEXT,
_REGEX_NEGATIVE_RESULTS), (ModelFileType.FILE_CONTENT, _REGEX_MODEL_FILE,
_NEGATIVE_TEXT, _REGEX_NEGATIVE_RESULTS),
(ModelFileType.FILE_NAME, _REGEX_MODEL_FILE, _POSITIVE_TEXT,
_REGEX_POSITIVE_RESULTS), (ModelFileType.FILE_CONTENT, _REGEX_MODEL_FILE,
_POSITIVE_TEXT, _REGEX_POSITIVE_RESULTS))
def test_classify(self, model_file_type, model_name, text,
expected_classification_result):
# Creates classifier.
model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, model_name))
if model_file_type is ModelFileType.FILE_NAME:
base_options = _BaseOptions(model_asset_path=model_path)
elif model_file_type is ModelFileType.FILE_CONTENT:
with open(model_path, 'rb') as f:
model_content = f.read()
base_options = _BaseOptions(model_asset_buffer=model_content)
else:
# Should never happen
raise ValueError('model_file_type is invalid.')
options = _TextClassifierOptions(base_options=base_options)
classifier = _TextClassifier.create_from_options(options)
# Performs text classification on the input.
text_result = classifier.classify(text)
# Comparing results.
self.assertTextClassifierResultEquals(
text_result, expected_classification_result
)
# Closes the classifier explicitly when the classifier is not used in
# a context.
classifier.close()
@parameterized.parameters((ModelFileType.FILE_NAME, _BERT_MODEL_FILE,
_NEGATIVE_TEXT, _BERT_NEGATIVE_RESULTS),
(ModelFileType.FILE_CONTENT, _BERT_MODEL_FILE,
_NEGATIVE_TEXT, _BERT_NEGATIVE_RESULTS))
def test_classify_in_context(self, model_file_type, model_name, text,
expected_classification_result):
# Creates classifier.
model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, model_name))
if model_file_type is ModelFileType.FILE_NAME:
base_options = _BaseOptions(model_asset_path=model_path)
elif model_file_type is ModelFileType.FILE_CONTENT:
with open(model_path, 'rb') as f:
model_content = f.read()
base_options = _BaseOptions(model_asset_buffer=model_content)
else:
# Should never happen
raise ValueError('model_file_type is invalid.')
options = _TextClassifierOptions(base_options=base_options)
with _TextClassifier.create_from_options(options) as classifier:
# Performs text classification on the input.
text_result = classifier.classify(text)
# Comparing results.
self.assertTextClassifierResultEquals(
text_result, expected_classification_result
)
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/text/text_embedder_test.py
+328
View File
@@ -0,0 +1,328 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.
"""Tests for text embedder."""
import enum
import os
from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from mediapipe.tasks.python.components.containers import embedding_result as embedding_result_module
from mediapipe.tasks.python.core import base_options as base_options_module
from mediapipe.tasks.python.test import test_utils
from mediapipe.tasks.python.text import text_embedder
_BaseOptions = base_options_module.BaseOptions
_Embedding = embedding_result_module.Embedding
_TextEmbedder = text_embedder.TextEmbedder
_TextEmbedderOptions = text_embedder.TextEmbedderOptions
_BERT_MODEL_FILE = 'mobilebert_embedding_with_metadata.tflite'
_REGEX_MODEL_FILE = 'regex_one_embedding_with_metadata.tflite'
_USE_MODEL_FILE = 'universal_sentence_encoder_qa_with_metadata.tflite'
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
# Tolerance for embedding vector coordinate values.
_EPSILON = 1e-4
# Tolerance for cosine similarity evaluation.
_SIMILARITY_TOLERANCE = 1e-3
class ModelFileType(enum.Enum):
FILE_CONTENT = 1
FILE_NAME = 2
class TextEmbedderTest(parameterized.TestCase):
def setUp(self):
super().setUp()
self.model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, _BERT_MODEL_FILE))
def test_create_from_file_succeeds_with_valid_model_path(self):
# Creates with default option and valid model file successfully.
with _TextEmbedder.create_from_model_path(self.model_path) as embedder:
self.assertIsInstance(embedder, _TextEmbedder)
def test_create_from_options_succeeds_with_valid_model_path(self):
# Creates with options containing model file successfully.
base_options = _BaseOptions(model_asset_path=self.model_path)
options = _TextEmbedderOptions(base_options=base_options)
with _TextEmbedder.create_from_options(options) as embedder:
self.assertIsInstance(embedder, _TextEmbedder)
def test_create_from_options_fails_with_invalid_model_path(self):
with self.assertRaisesRegex(
FileNotFoundError,
'Unable to open file at /path/to/invalid/model.tflite',
):
base_options = _BaseOptions(
model_asset_path='/path/to/invalid/model.tflite')
options = _TextEmbedderOptions(base_options=base_options)
_TextEmbedder.create_from_options(options)
def test_create_from_options_succeeds_with_valid_model_content(self):
# Creates with options containing model content successfully.
with open(self.model_path, 'rb') as f:
base_options = _BaseOptions(model_asset_buffer=f.read())
options = _TextEmbedderOptions(base_options=base_options)
embedder = _TextEmbedder.create_from_options(options)
self.assertIsInstance(embedder, _TextEmbedder)
def _check_embedding_value(self, result, expected_first_value):
# Check embedding first value.
self.assertAlmostEqual(
result.embeddings[0].embedding[0], expected_first_value, delta=_EPSILON)
def _check_embedding_size(self, result, quantize, expected_embedding_size):
# Check embedding size.
self.assertLen(result.embeddings, 1)
embedding_result = result.embeddings[0]
self.assertLen(embedding_result.embedding, expected_embedding_size)
if quantize:
self.assertEqual(embedding_result.embedding.dtype, np.uint8)
else:
self.assertEqual(embedding_result.embedding.dtype, float)
def _check_cosine_similarity(self, result0, result1, expected_similarity):
# Checks cosine similarity.
similarity = _TextEmbedder.cosine_similarity(result0.embeddings[0],
result1.embeddings[0])
self.assertAlmostEqual(
similarity, expected_similarity, delta=_SIMILARITY_TOLERANCE)
@parameterized.parameters(
(
False,
False,
_BERT_MODEL_FILE,
ModelFileType.FILE_NAME,
0.9624276,∂ç
512,
(21.2054, 19.6843),
),
(
True,
False,
_BERT_MODEL_FILE,
ModelFileType.FILE_NAME,
0.9624276,
512,
(0.062578, 0.067393),
),
(
False,
False,
_REGEX_MODEL_FILE,
ModelFileType.FILE_NAME,
0.999937,
16,
(0.0309356, 0.0312863),
),
(
True,
False,
_REGEX_MODEL_FILE,
ModelFileType.FILE_CONTENT,
0.999937,
16,
(0.549632, 0.552879),
),
(
False,
False,
_USE_MODEL_FILE,
ModelFileType.FILE_NAME,
0.851961,
100,
(1.422951, 1.404664),
),
(
True,
False,
_USE_MODEL_FILE,
ModelFileType.FILE_CONTENT,
0.851961,
100,
(0.127049, 0.125416),
),
)
def test_embed(self, l2_normalize, quantize, model_name, model_file_type,
expected_similarity, expected_size, expected_first_values):
# Creates embedder.
model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, model_name))
if model_file_type is ModelFileType.FILE_NAME:
base_options = _BaseOptions(model_asset_path=model_path)
elif model_file_type is ModelFileType.FILE_CONTENT:
with open(model_path, 'rb') as f:
model_content = f.read()
base_options = _BaseOptions(model_asset_buffer=model_content)
else:
# Should never happen
raise ValueError('model_file_type is invalid.')
options = _TextEmbedderOptions(
base_options=base_options, l2_normalize=l2_normalize, quantize=quantize)
embedder = _TextEmbedder.create_from_options(options)
# Extracts both embeddings.
positive_text0 = "it's a charming and often affecting journey"
positive_text1 = 'what a great and fantastic trip'
result0 = embedder.embed(positive_text0)
result1 = embedder.embed(positive_text1)
# Checks embeddings and cosine similarity.
expected_result0_value, expected_result1_value = expected_first_values
self._check_embedding_size(result0, quantize, expected_size)
self._check_embedding_size(result1, quantize, expected_size)
self._check_embedding_value(result0, expected_result0_value)
self._check_embedding_value(result1, expected_result1_value)
self._check_cosine_similarity(result0, result1, expected_similarity)
# Closes the embedder explicitly when the embedder is not used in
# a context.
embedder.close()
@parameterized.parameters(
(
False,
False,
_BERT_MODEL_FILE,
ModelFileType.FILE_NAME,
0.962427,
512,
(21.2054, 19.684337),
),
(
True,
False,
_BERT_MODEL_FILE,
ModelFileType.FILE_NAME,
0.962427,
512,
(0.0625787, 0.0673937),
),
(
False,
False,
_REGEX_MODEL_FILE,
ModelFileType.FILE_NAME,
0.999937,
16,
(0.0309356, 0.0312863),
),
(
True,
False,
_REGEX_MODEL_FILE,
ModelFileType.FILE_CONTENT,
0.999937,
16,
(0.549632, 0.552879),
),
(
False,
False,
_USE_MODEL_FILE,
ModelFileType.FILE_NAME,
0.851961,
100,
(1.422951, 1.404664),
),
(
True,
False,
_USE_MODEL_FILE,
ModelFileType.FILE_CONTENT,
0.851961,
100,
(0.127049, 0.125416),
),
)
def test_embed_in_context(self, l2_normalize, quantize, model_name,
model_file_type, expected_similarity, expected_size,
expected_first_values):
# Creates embedder.
model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, model_name))
if model_file_type is ModelFileType.FILE_NAME:
base_options = _BaseOptions(model_asset_path=model_path)
elif model_file_type is ModelFileType.FILE_CONTENT:
with open(model_path, 'rb') as f:
model_content = f.read()
base_options = _BaseOptions(model_asset_buffer=model_content)
else:
# Should never happen
raise ValueError('model_file_type is invalid.')
options = _TextEmbedderOptions(
base_options=base_options, l2_normalize=l2_normalize, quantize=quantize)
with _TextEmbedder.create_from_options(options) as embedder:
# Extracts both embeddings.
positive_text0 = "it's a charming and often affecting journey"
positive_text1 = 'what a great and fantastic trip'
result0 = embedder.embed(positive_text0)
result1 = embedder.embed(positive_text1)
# Checks embeddings and cosine similarity.
expected_result0_value, expected_result1_value = expected_first_values
self._check_embedding_size(result0, quantize, expected_size)
self._check_embedding_size(result1, quantize, expected_size)
self._check_embedding_value(result0, expected_result0_value)
self._check_embedding_value(result1, expected_result1_value)
self._check_cosine_similarity(result0, result1, expected_similarity)
@parameterized.parameters(
# TODO: The similarity should likely be lower
(_BERT_MODEL_FILE, 0.98103),
(_USE_MODEL_FILE, 0.780334),
)
def test_embed_with_different_themes(self, model_file, expected_similarity):
# Creates embedder.
model_path = test_utils.get_test_data_path(
os.path.join(_TEST_DATA_DIR, model_file)
)
base_options = _BaseOptions(model_asset_path=model_path)
options = _TextEmbedderOptions(base_options=base_options)
embedder = _TextEmbedder.create_from_options(options)
# Extracts both embeddings.
text0 = (
'When you go to this restaurant, they hold the pancake upside-down '
"before they hand it to you. It's a great gimmick."
)
result0 = embedder.embed(text0)
text1 = "Let's make a plan to steal the declaration of independence."
result1 = embedder.embed(text1)
similarity = _TextEmbedder.cosine_similarity(
result0.embeddings[0], result1.embeddings[0]
)
self.assertAlmostEqual(
similarity, expected_similarity, delta=_SIMILARITY_TOLERANCE
)
# Closes the embedder explicitly when the embedder is not used in
# a context.
embedder.close()
if __name__ == '__main__':
absltest.main()
backend/.venv/Lib/site-packages/mediapipe/tasks/python/test/vision/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The MediaPipe Authors.
#
# 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.

Some files were not shown because too many files have changed in this diff Show More