Python coremltools.proto方法代码示例

# 需要导入模块: import coremltools [as 别名]
# 或者: from coremltools import proto [as 别名]
def __init__(self, model):
        from coremltools.proto import Model_pb2

        # sanity check.
        if not isinstance(model, Model_pb2.Model):
            raise TypeError("Coreml layer of type %s is not supported." % type(model))
        super(CoremlGraph, self).__init__(model)
        self.model = model 

# 需要导入模块: import coremltools [as 别名]
# 或者: from coremltools import proto [as 别名]
def convert_multiarray_output_to_image(spec, feature_name, is_bgr=False):  
    """  
    Convert an output multiarray to be represented as an image  
    This will modify the Model_pb spec passed in.  
    Example:  
        model = coremltools.models.MLModel('MyNeuralNetwork.mlmodel')  
        spec = model.get_spec()  
        convert_multiarray_output_to_image(spec,'imageOutput',is_bgr=False)  
        newModel = coremltools.models.MLModel(spec)  
        newModel.save('MyNeuralNetworkWithImageOutput.mlmodel')  
    Parameters  
    ----------  
    spec: Model_pb  
        The specification containing the output feature to convert  
    feature_name: str  
        The name of the multiarray output feature you want to convert  
    is_bgr: boolean  
        If multiarray has 3 channels, set to True for RGB pixel order or false for BGR  
    """  
    for output in spec.description.output:  
        if output.name != feature_name:  
            continue  
        if output.type.WhichOneof('Type') != 'multiArrayType':  
            raise ValueError("%s is not a multiarray type" % output.name)  
        array_shape = tuple(output.type.multiArrayType.shape)  
        channels, height, width = array_shape  
        from coremltools.proto import FeatureTypes_pb2 as ft  
        if channels == 1:  
            output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('GRAYSCALE')  
        elif channels == 3:  
            if is_bgr:  
                output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('BGR')  
            else:  
                output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('RGB')  
        else:  
            raise ValueError("Channel Value %d not supported for image inputs" % channels)  
        output.type.imageType.width = width  
        output.type.imageType.height = height 

# 需要导入模块: import coremltools [as 别名]
# 或者: from coremltools import proto [as 别名]
def _convert_k_neighbors_classifier(model, input_name, output_name):
    """Convert the scikit KNeighborsClassifier to CoreML. Assumes initial validation of the scikit model has been done."""

    spec = coremltools.proto.Model_pb2.Model()
    spec.specificationVersion = coremltools.SPECIFICATION_VERSION

    spec.kNearestNeighborsClassifier.numberOfNeighbors.defaultValue = model.n_neighbors
    spec.kNearestNeighborsClassifier.numberOfNeighbors.range.minValue = 1
    spec.kNearestNeighborsClassifier.numberOfNeighbors.range.maxValue = _number_of_samples(
        model, spec
    )  # is there a better heuristic to use here?

    number_of_dimensions = 0
    if _is_algorithm_brute(model):
        number_of_dimensions = model._fit_X.shape[1]
        spec.kNearestNeighborsClassifier.nearestNeighborsIndex.linearIndex.MergeFromString(
            b""
        )
    elif _is_algorithm_kd_tree(model):
        npdata = np.asarray(model._tree.data)
        number_of_dimensions = get_input_dimension(model)
        spec.kNearestNeighborsClassifier.nearestNeighborsIndex.singleKdTreeIndex.leafSize = (
            model.leaf_size
        )
    else:
        raise TypeError(
            "KNeighbors algorithm not supported for CoreML conversion: {}".format(
                model.algorithm
            )
        )
    spec.kNearestNeighborsClassifier.nearestNeighborsIndex.numberOfDimensions = (
        number_of_dimensions
    )

    # Make sure the distance function is set
    spec.kNearestNeighborsClassifier.nearestNeighborsIndex.squaredEuclideanDistance.MergeFromString(
        b""
    )

    input_features = spec.description.input.add()
    input_features.name = input_name[0][0]
    input_features.type.multiArrayType.shape.extend([number_of_dimensions])
    input_features.type.multiArrayType.dataType = (
        FeatureTypes_pb2.ArrayFeatureType.FLOAT32
    )

    output_label = spec.description.output.add()
    output_label.name = output_name[0][0]

    # predictedFeatureName is required since KNN is a classifier and it should be same as outputName.
    spec.description.predictedFeatureName = output_label.name

    # Need to confirm if scikit only accepts integer labels
    output_label.type.int64Type.MergeFromString(b"")
    spec.kNearestNeighborsClassifier.uniformWeighting.MergeFromString(b"")

    _extract_training_data(model, spec)

    return spec