#import "CoreMLPythonArray.h"

#import "CoreMLPythonUtils.h"

#include <pybind11/eval.h>

#include <pybind11/numpy.h>

#include <iomanip> // for std::setfill etc

#import <Accelerate/Accelerate.h>

#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <numpy/arrayobject.h>

#define PyAnyInteger_Check(name) (PyLong_Check(name) || (_import_array(), PyArray_IsScalar(name, Integer)))

using namespace CoreML::Python;

bool Utils::isCompiledModelPath(const std::string& path) {

const std::string fileExtension = ".mlmodelc";

size_t start = path.length() - fileExtension.length();

if (path.back() == '/') {

start--;

}

const std::string match = path.substr(start, fileExtension.length());

return (match == fileExtension);

}

NSURL * Utils::stringToNSURL(const std::string& str) {

NSString *nsstr = [NSString stringWithUTF8String:str.c_str()];

return [NSURL fileURLWithPath:nsstr];

}

void Utils::handleError(NSError *error) {

if (error != nil) {

NSString *formatted = [NSString stringWithFormat:@"%@", [error userInfo]];

throw std::runtime_error([formatted UTF8String]);

}

}

MLDictionaryFeatureProvider * Utils::dictToFeatures(const py::dict& dict, NSError * __autoreleasing *error) {

NSError *localError;

MLDictionaryFeatureProvider * feautreProvider;

@autoreleasepool {

NSMutableDictionary<NSString *, NSObject *> *inputDict = [[NSMutableDictionary<NSString *, NSObject *> alloc] init];

for (const auto element : dict) {

std::string key = element.first.cast<std::string>();

NSString *nsKey = [NSString stringWithUTF8String:key.c_str()];

id nsValue = Utils::convertValueToObjC(element.second);

inputDict[nsKey] = nsValue;

}

feautreProvider = [[MLDictionaryFeatureProvider alloc] initWithDictionary:inputDict error:&localError];

}

if (error != NULL) {

*error = localError;

}

return feautreProvider;

}

py::dict Utils::featuresToDict(id<MLFeatureProvider> features) {

@autoreleasepool {

py::dict ret;

NSSet<NSString *> *keys = [features featureNames];

for (NSString *key in keys) {

MLFeatureValue *value = [features featureValueForName:key];

py::str pyKey = py::str([key UTF8String]);

py::object pyValue = convertValueToPython(value);

ret[pyKey] = pyValue;

}

return ret;

}

}

template<typename KEYTYPE>

static NSObject * convertDictKey(const KEYTYPE& k);

template<>

NSObject * convertDictKey(const int64_t& k) {

return [NSNumber numberWithLongLong:k];

}

template<>

NSObject * convertDictKey(const std::string& k) {

return [NSString stringWithUTF8String:k.c_str()];

}

template<typename VALUETYPE>

static NSNumber * convertDictValue(const VALUETYPE& v);

template<>

NSNumber * convertDictValue(const int64_t& v) {

return [NSNumber numberWithLongLong:v];

}

template<>

NSNumber * convertDictValue(const double& v) {

return [NSNumber numberWithDouble:v];

}

template<typename KEYTYPE, typename VALUETYPE>

static MLFeatureValue * convertToNSDictionary(const std::unordered_map<KEYTYPE, VALUETYPE>& dict) {

NSMutableDictionary<NSObject *, NSNumber *> *nsDict = [[NSMutableDictionary<NSObject *, NSNumber *> alloc] init];

for (const auto& pair : dict) {

NSObject<NSCopying> *key = (NSObject<NSCopying> *)convertDictKey(pair.first);

NSNumber *value = convertDictValue(pair.second);

assert(key != nil);

nsDict[key] = value;

}

NSError *error = nil;

MLFeatureValue * ret = [MLFeatureValue featureValueWithDictionary:nsDict error:&error];

if (error != nil) {

throw std::runtime_error(error.localizedDescription.UTF8String);

}

return ret;

}

static MLFeatureValue * convertValueToDictionary(const py::handle& handle) {

if(!PyDict_Check(handle.ptr())) {

throw std::runtime_error("Not a dictionary.");

}

// Get the first value in the dictionary; use that as a hint.

PyObject *key = nullptr, *value = nullptr;

Py_ssize_t pos = 0;

int has_values = PyDict_Next(handle.ptr(), &pos, &key, &value);

// Is it an empty dict? If so, just return an empty dictionary.

if(!has_values) {

return [MLFeatureValue featureValueWithDictionary:@{} error:nullptr];

}

if(PyAnyInteger_Check(key)) {

if(PyAnyInteger_Check(value)) {

auto dict = handle.cast<std::unordered_map<int64_t, int64_t> >();

return convertToNSDictionary(dict);

} else if(PyFloat_Check(value)) {

auto dict = handle.cast<std::unordered_map<int64_t, double> >();

return convertToNSDictionary(dict);

} else {

throw std::runtime_error("Unknown value type for int key in dictionary.");

}

} else if (PyBytes_Check(key) || PyUnicode_Check(key)) {

if(PyAnyInteger_Check(value)) {

auto dict = handle.cast<std::unordered_map<std::string, int64_t> >();

return convertToNSDictionary(dict);

} else if(PyFloat_Check(value)) {

auto dict = handle.cast<std::unordered_map<std::string, double> >();

return convertToNSDictionary(dict);

} else {

throw std::runtime_error("Invalid value type for string key in dictionary.");

}

} else {

throw std::runtime_error("Invalid key type dictionary.");

}

}

static MLFeatureValue * convertValueToArray(const py::handle& handle) {

// if this line throws, it can't be an array (caller should catch)

py::array buf = handle.cast<py::array>();

if (buf.shape() == nullptr) {

throw std::runtime_error("no shape, can't be an array");

}

PybindCompatibleArray *array = [[PybindCompatibleArray alloc] initWithArray:buf];

return [MLFeatureValue featureValueWithMultiArray:array];

}

static MLFeatureValue * convertValueToSequence(const py::handle& handle) {

#pragma unused (handle)

/*

return nil;

}

static void handleCVReturn(CVReturn status) {

if (status != kCVReturnSuccess) {

std::stringstream msg;

msg << "Got unexpected return code " << status << " from CoreVideo.";

throw std::runtime_error(msg.str());

}

}

static MLFeatureValue * convertValueToImage(const py::handle& handle) {

// assumes handle is a valid PIL image!

CVPixelBufferRef pixelBuffer = nil;

size_t width = handle.attr("width").cast<size_t>();

size_t height = handle.attr("height").cast<size_t>();

OSType format;

std::string formatStr = handle.attr("mode").cast<std::string>();

if (formatStr == "RGB") {

format = kCVPixelFormatType_32BGRA;

} else if (formatStr == "RGBA") {

format = kCVPixelFormatType_32BGRA;

} else if (formatStr == "L") {

format = kCVPixelFormatType_OneComponent8;

} else if (formatStr == "F") {

format = kCVPixelFormatType_OneComponent16Half;

} else {

std::stringstream msg;

msg << "Unsupported image type " << formatStr << ". ";

msg << "Supported types are: RGB, RGBA, L.";

throw std::runtime_error(msg.str());

}

CVReturn status = CVPixelBufferCreate(kCFAllocatorDefault, width, height, format, NULL, &pixelBuffer);

handleCVReturn(status);

// get bytes out of the PIL image

py::object tobytes = handle.attr("tobytes");

py::object bytesResult = tobytes();

assert(PyBytes_Check(bytesResult.ptr()));

Py_ssize_t bytesLength = PyBytes_Size(bytesResult.ptr());

assert(bytesLength >= 0);

const char *bytesPtr = PyBytes_AsString(bytesResult.ptr());

// copy data into the CVPixelBuffer

status = CVPixelBufferLockBaseAddress(pixelBuffer, 0);

handleCVReturn(status);

void *baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer);

assert(baseAddress != nullptr);

assert(!CVPixelBufferIsPlanar(pixelBuffer));

size_t bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer);

const char *srcPointer = bytesPtr;

vImage_Buffer srcBuffer;

memset(&srcBuffer, 0, sizeof(srcBuffer));

srcBuffer.data = const_cast<char *>(srcPointer);

srcBuffer.width = width;

srcBuffer.height = height;

vImage_Buffer dstBuffer;

memset(&dstBuffer, 0, sizeof(dstBuffer));

dstBuffer.data = baseAddress;

dstBuffer.width = width;

dstBuffer.height = height;

if (formatStr == "RGB") {

// convert RGB to BGRA

assert(bytesLength == width * height * 3);

srcBuffer.rowBytes = width * 3;

dstBuffer.rowBytes = bytesPerRow;

vImageConvert_RGB888toBGRA8888(&srcBuffer, NULL, 255, &dstBuffer, false, 0);

} else if (formatStr == "RGBA") {

// convert RGBA to BGRA

assert(bytesLength == width * height * 4);

srcBuffer.rowBytes = width * 4;

dstBuffer.rowBytes = bytesPerRow;

uint8_t permuteMap[4] = { 2, 1, 0, 3 };

vImagePermuteChannels_ARGB8888(&srcBuffer, &dstBuffer, permuteMap, 0);

} else if (formatStr == "L") {

// 8 bit grayscale.

assert(bytesLength == width * height);

srcBuffer.rowBytes = width;

dstBuffer.rowBytes = bytesPerRow;

vImageCopyBuffer(&srcBuffer, &dstBuffer, 1, 0);

} else if (formatStr == "F") {

// convert Float32 to Float16.

assert(bytesLength == width * height * sizeof(Float32));

srcBuffer.rowBytes = width * sizeof(Float32);

dstBuffer.rowBytes = bytesPerRow;

vImageConvert_PlanarFtoPlanar16F(&srcBuffer, &dstBuffer, 0);

} else {

std::stringstream msg;

msg << "Unsupported image type " << formatStr << ". ";

msg << "Supported types are: RGB, RGBA, L.";

throw std::runtime_error(msg.str());

}

#ifdef COREML_SHOW_PIL_IMAGES

if (formatStr == "RGB") {

// for debugging purposes, convert back to PIL image and show it

py::object scope = py::module::import("__main__").attr("__dict__");

py::eval<py::eval_single_statement>("import PIL.Image", scope);

py::object pilImage = py::eval<py::eval_expr>("PIL.Image");

std::string cvPixelStr(count, 0);

const char *basePtr = static_cast<char *>(baseAddress);

for (size_t row = 0; row < height; row++) {

for (size_t col = 0; col < width; col++) {

for (size_t color = 0; color < 3; color++) {

cvPixelStr[(row * width * 3) + (col*3) + color] = basePtr[(row * bytesPerRow) + (col*4) + color + 1];

}

}

}

py::bytes cvPixelBytes = py::bytes(cvPixelStr);

py::object frombytes = pilImage.attr("frombytes");

py::str mode = "RGB";

auto size = py::make_tuple(width, height);

py::object img = frombytes(mode, size, cvPixelBytes);

img.attr("show")();

}

#endif

status = CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);

handleCVReturn(status);

MLFeatureValue *fv = [MLFeatureValue featureValueWithPixelBuffer:pixelBuffer];

CVPixelBufferRelease(pixelBuffer);

return fv;

}

static bool IsPILImage(const py::handle& handle) {

// TODO put try/catch around this?

try {

py::module::import("PIL.Image");

} catch(...) {

return false;

}

py::object scope = py::module::import("__main__").attr("__dict__");

py::eval<py::eval_single_statement>("import PIL.Image", scope);

py::handle imageTypeHandle = py::eval<py::eval_expr>("PIL.Image.Image", scope);

assert(PyType_Check(imageTypeHandle.ptr())); // should be a Python type

return PyObject_TypeCheck(handle.ptr(), (PyTypeObject *)(imageTypeHandle.ptr()));

}

MLFeatureValue * Utils::convertValueToObjC(const py::handle& handle) {

if (PyAnyInteger_Check(handle.ptr())) {

try {

int64_t val = handle.cast<int64_t>();

return [MLFeatureValue featureValueWithInt64:val];

} catch(...) {}

}

if (PyFloat_Check(handle.ptr())) {

try {

double val = handle.cast<double>();

return [MLFeatureValue featureValueWithDouble:val];

} catch(...) {}

}

if (PyBytes_Check(handle.ptr()) || PyUnicode_Check(handle.ptr())) {

try {

std::string val = handle.cast<std::string>();

return [MLFeatureValue featureValueWithString:[NSString stringWithUTF8String:val.c_str()]];

} catch(...) {}

}

if (PyDict_Check(handle.ptr())) {

try {

return convertValueToDictionary(handle);

} catch(...) {}

}

if(PyList_Check(handle.ptr()) || PyTuple_Check(handle.ptr())) {

try {

return convertValueToSequence(handle);

} catch(...) {}

}

if(PyObject_CheckBuffer(handle.ptr())) {

try {

return convertValueToArray(handle);

} catch(...) {}

}

if (IsPILImage(handle)) {

return convertValueToImage(handle);

}

py::print("Error: value type not convertible:");

py::print(handle);

throw std::runtime_error("value type not convertible");

}

std::vector<size_t> Utils::convertNSArrayToCpp(NSArray<NSNumber *> *array) {

std::vector<size_t> ret;

for (NSNumber *value in array) {

ret.push_back(value.unsignedLongValue);

}

return ret;

}

NSArray<NSNumber *>* Utils::convertCppArrayToObjC(const std::vector<size_t>& array) {

NSMutableArray<NSNumber *>* ret = [[NSMutableArray<NSNumber *> alloc] init];

for (size_t element : array) {

[ret addObject:[NSNumber numberWithUnsignedLongLong:element]];

}

return ret;

}

static size_t sizeOfArrayElement(MLMultiArrayDataType type) {

switch (type) {

case MLMultiArrayDataTypeInt32:

return sizeof(int32_t);

case MLMultiArrayDataTypeFloat32:

case MLMultiArrayDataTypeFloat16:

return sizeof(float);

case MLMultiArrayDataTypeDouble:

return sizeof(double);

default:

assert(false);

return sizeof(double);

}

}

py::object Utils::convertArrayValueToPython(MLMultiArray *value) {

if (value == nil) {

return py::none();

}

MLMultiArrayDataType type = value.dataType;

if (type == MLMultiArrayDataTypeFloat16) {

// Cast to fp32 because py:array doesn't support fp16.

// TODO: rdar://92239209 : return np.float16 instead of np.float32 when multiarray type is Float16

value = [MLMultiArray multiArrayByConcatenatingMultiArrays:@[value] alongAxis:0 dataType:MLMultiArrayDataTypeFloat32];

type = value.dataType;

}

std::vector<size_t> shape = Utils::convertNSArrayToCpp(value.shape);

std::vector<size_t> strides = Utils::convertNSArrayToCpp(value.strides);

// convert strides to numpy (bytes) instead of mlkit (elements)

for (size_t& stride : strides) {

stride *= sizeOfArrayElement(type);

}

__block py::object array;

[value getBytesWithHandler:^(const void *bytes, NSInteger size) {

switch (type) {

#if BUILT_WITH_MACOS26_SDK

case MLMultiArrayDataTypeInt8:

array = py::array(shape, strides, reinterpret_cast<const int8_t *>(bytes));

break;

#endif

case MLMultiArrayDataTypeInt32:

array = py::array(shape, strides, reinterpret_cast<const int32_t *>(bytes));

break;

case MLMultiArrayDataTypeFloat32:

array = py::array(shape, strides, reinterpret_cast<const float *>(bytes));

break;

case MLMultiArrayDataTypeFloat64:

array = py::array(shape, strides, reinterpret_cast<const double *>(bytes));

break;

default:

assert(false);

array = py::object();

}

}];

return array;

}

py::object Utils::convertDictionaryValueToPython(NSDictionary<NSObject *,NSNumber *> * dict) {

if (dict == nil) {

return py::none();

}

py::dict ret;

for (NSObject *key in dict) {

py::object pykey;

if ([key isKindOfClass:[NSNumber class]]) {

// can assume int32_t -- we only allow arrays of int or string keys

NSNumber *nskey = static_cast<NSNumber *>(key);

pykey = py::int_([nskey integerValue]);

} else {

assert([key isKindOfClass:[NSString class]]);

NSString *nskey = static_cast<NSString *>(key);

pykey = py::str([nskey UTF8String]);

}

NSNumber *value = dict[key];

ret[pykey] = py::float_([value doubleValue]);

}

return std::move(ret);

}

py::object Utils::convertImageValueToPython(CVPixelBufferRef value) {

if (CVPixelBufferIsPlanar(value)) {

throw std::runtime_error("Only non-planar CVPixelBuffers are currently supported by this Python binding.");

}

// supports grayscale and BGRA format types

auto formatType = CVPixelBufferGetPixelFormatType(value);

assert(formatType == kCVPixelFormatType_32BGRA

|| formatType == kCVPixelFormatType_OneComponent8

|| formatType == kCVPixelFormatType_OneComponent16Half);

auto height = CVPixelBufferGetHeight(value);

auto width = CVPixelBufferGetWidth(value);

py::str mode;

size_t dstBytesPerRow = 0;

if (formatType == kCVPixelFormatType_32BGRA) {

mode = "RGBA";

dstBytesPerRow = width * 4;

} else if (formatType == kCVPixelFormatType_OneComponent8) {

mode = "L";

dstBytesPerRow = width * sizeof(uint8_t);

} else if (formatType == kCVPixelFormatType_OneComponent16Half) {

mode = "F";

dstBytesPerRow = width * sizeof(Float32);

} else {

std::stringstream msg;

msg << "Unsupported pixel format type: " << std::hex << std::setfill('0') << std::setw(4) << formatType << ". ";

throw std::runtime_error(msg.str());

}

PyObject *dstPyBytes = PyBytes_FromStringAndSize(NULL, height * dstBytesPerRow);

if (!dstPyBytes) {

throw std::bad_alloc();

}

auto result = CVPixelBufferLockBaseAddress(value, kCVPixelBufferLock_ReadOnly);

assert(result == kCVReturnSuccess);

uint8_t *src = reinterpret_cast<uint8_t*>(CVPixelBufferGetBaseAddress(value));

assert(src != nullptr);

size_t srcBytesPerRow = CVPixelBufferGetBytesPerRow(value);

// Prepare for vImage blitting

vImage_Buffer srcBuffer;

memset(&srcBuffer, 0, sizeof(srcBuffer));

srcBuffer.data = src;

srcBuffer.width = width;

srcBuffer.height = height;

srcBuffer.rowBytes = srcBytesPerRow;

vImage_Buffer dstBuffer;

memset(&dstBuffer, 0, sizeof(dstBuffer));

dstBuffer.data = PyBytes_AS_STRING(dstPyBytes);

dstBuffer.width = width;

dstBuffer.height = height;

dstBuffer.rowBytes = dstBytesPerRow;

if (formatType == kCVPixelFormatType_32BGRA) {

// convert BGRA to RGBA

uint8_t permuteMap[4] = { 2, 1, 0, 3 };

vImagePermuteChannels_ARGB8888(&srcBuffer, &dstBuffer, permuteMap, 0);

} else if (formatType == kCVPixelFormatType_OneComponent8) {

vImageCopyBuffer(&srcBuffer, &dstBuffer, 1, 0);

} else if (formatType == kCVPixelFormatType_OneComponent16Half) {

vImageConvert_Planar16FtoPlanarF(&srcBuffer, &dstBuffer, 0);

} else {

std::stringstream msg;

msg << "Unsupported pixel format type: " << std::hex << std::setfill('0') << std::setw(4) << formatType << ". ";

throw std::runtime_error(msg.str());

}

result = CVPixelBufferUnlockBaseAddress(value, kCVPixelBufferLock_ReadOnly);

assert(result == kCVReturnSuccess);

py::object scope = py::module::import("__main__").attr("__dict__");

py::eval<py::eval_single_statement>("import PIL.Image", scope);

py::object pilImage = py::eval<py::eval_expr>("PIL.Image", scope);

py::object frombytes = pilImage.attr("frombytes");

py::bytes dstBytes = py::reinterpret_steal<py::bytes>(dstPyBytes); // transfer ownership of `dstPyBytes` to `dstBytes`

py::object img = frombytes(mode, py::make_tuple(width, height), dstBytes);

return img;

}

py::object Utils::convertSequenceValueToPython(MLSequence *seq) {

if (seq == nil) {

return py::none();

}

py::list ret;

if(seq.type == MLFeatureTypeString) {

for(NSString* s in seq.stringValues) {

ret.append(py::str(s.UTF8String));

}

} else if(seq.type == MLFeatureTypeInt64) {

for(NSNumber* n in seq.stringValues) {

ret.append(py::int_(n.longLongValue) );

}

} else {

throw std::runtime_error("Error: Unrecognized sequence type.");

}

return std::move(ret);

}

py::object Utils::convertValueToPython(MLFeatureValue *value) {

switch ([value type]) {

case MLFeatureTypeInt64:

return py::int_(value.int64Value);

case MLFeatureTypeMultiArray:

return convertArrayValueToPython(value.multiArrayValue);

case MLFeatureTypeImage:

return convertImageValueToPython(value.imageBufferValue);

case MLFeatureTypeDouble:

return py::float_(value.doubleValue);

case MLFeatureTypeString:

return py::str(value.stringValue.UTF8String);

case MLFeatureTypeDictionary:

return convertDictionaryValueToPython(value.dictionaryValue);

case MLFeatureTypeSequence:

if (@available(macOS 10.14, *)) {

return convertSequenceValueToPython(value.sequenceValue);

} else {

// Fallback on earlier versions

}

case MLFeatureTypeInvalid:

assert(false);

return py::none();

}

return py::object();

}