Add PIL Image support for remote device inference

[soof-golan]

This PR:

Enables passing and receiving PIL.Image.Image objects to remote inference by:

• Adding UInt8 dtype support to TensorDescriptor
• Creating a CVPixelBuffer as the input based on the model's input metadata
• Retrieving the contents of a CVPixelBuffer based on the model's output metadata

Testing Approach

Manual testing was performed to ensure the code works as intended.
I'm not sure how to integrate this into the main repo in a clean way:

import asyncio
import sys

import tempfile
import numpy as np
from PIL import Image

import coremltools as ct
from coremltools.models.ml_program.experimental.remote_device import (
    Device,
    AppSigningCredentials,
    DeviceType,
)
from coremltools.models.ml_program.experimental.async_wrapper import (
    RemoteMLModelAsyncWrapper,
)

async def main():
    model_path = tempfile.mkdtemp("image_io_model.mlpackage")
    create_model(model_path)

    team_id = "TEAMID"  # Replace with your Apple Developer Team ID
    if team_id == "TEAMID":
        print("Please set your Apple Developer Team ID in the code.")
        return 1
    
    credentials = AppSigningCredentials(development_team=team_id)
    connected_device = Device.get_connected_devices(device_type=DeviceType.IPHONE)[0]
    print(f"   Connected to: {connected_device.name}")

    prepared_device = await connected_device.prepare_for_model_debugging(
        credentials=credentials,
        clean=True,
    )

    img_array = np.zeros((64, 64, 3), dtype=np.uint8)
    img_array[0:32, 0:32] = [255, 0, 0]  # Red
    img_array[0:32, 32:64] = [0, 255, 0]  # Green
    img_array[32:64, 0:32] = [0, 0, 255]  # Blue
    img_array[32:64, 32:64] = [255, 255, 0]  # Yellow
    pil_image = Image.fromarray(img_array, mode="RGB")

    wrapper = RemoteMLModelAsyncWrapper(
        spec_or_path=model_path,
        weights_dir="",
        device=prepared_device,
        compute_units=ct.ComputeUnit.ALL,
    )
    await wrapper.load()

    outputs = await wrapper.predict({"image": pil_image})
    output_array = outputs["output_image"]

    errors = []
    out = output_array.astype(np.float32)

    if not np.allclose(out[0, 0], [255, 0, 0], atol=10):
        errors.append(f"Red wrong: {out[0, 0]}")
    if not np.allclose(out[0, 32], [0, 255, 0], atol=10):
        errors.append(f"Green wrong: {out[0, 32]}")
    if not np.allclose(out[32, 0], [0, 0, 255], atol=10):
        errors.append(f"Blue wrong: {out[32, 0]}")
    if not np.allclose(out[32, 32], [255, 255, 0], atol=10):
        errors.append(f"Yellow wrong: {out[32, 32]}")

    if errors:
        print("   ERRORS:")
        for e in errors:
            print(f"   - {e}")
        result = 1
    else:
        print("   All RGB values correct!")
        result = 0

    await wrapper.unload()

    print("SUCCESS!" if result == 0 else "FAILURE!")
    return result


def create_model(model_path):
    import coremltools as ct
    from coremltools.models.neural_network import NeuralNetworkBuilder

    input_features = [("image", ct.models.datatypes.Array(3, 64, 64))]
    output_features = [("output_image", ct.models.datatypes.Array(3, 64, 64))]

    builder = NeuralNetworkBuilder(input_features, output_features)
    builder.add_activation("identity", "LINEAR", "image", "output_image", [1.0, 0.0])

    builder.spec.description.input[0].type.imageType.width = 64
    builder.spec.description.input[0].type.imageType.height = 64
    builder.spec.description.input[
        0
    ].type.imageType.colorSpace = ct.proto.FeatureTypes_pb2.ImageFeatureType.RGB

    builder.spec.description.output[0].type.imageType.width = 64
    builder.spec.description.output[0].type.imageType.height = 64
    builder.spec.description.output[
        0
    ].type.imageType.colorSpace = ct.proto.FeatureTypes_pb2.ImageFeatureType.RGB

    model = ct.models.MLModel(builder.spec, skip_model_load=True)
    model.save(model_path)

if __name__ == "__main__":
    sys.exit(asyncio.run(main()))

[cymbalrush]

I have a couple of minor comments, but overall this looks very good—thank you! May I ask how you’re currently using this feature? Your input will help us enhance it.

[cymbalrush]

nit: Could we please make it a named tuple (data: [Data], shape: [Int], strides: [Int])?

[cymbalrush]

nit: Instead of force unwrapping baseAddress, could we do a guard and throw an error instead?

[TobyRoseman]

CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/2295212603