CoreML Model spec - change output type to dictionary [Double : String]

Hello everybody,

For the past week I have been struggling to run inference on a classifier I built using Google's AutoML Vision tool.

At first I thought everything would go smoothly because Google allows to export a CoreML version of the final model. I assumed I would only need to use Apple's CoreML library to make it work. When I export the model Google provides a .mlmodel file and a dict.txt file with the classification labels. For the current model I have 100 labels.

This is my Swift code to run inference on the model.

Code Block 
private lazy var classificationRequest: VNCoreMLRequest = {
        do {
            let classificationModel = try VNCoreMLModel(for: NewGenusModel().model)
            let request = VNCoreMLRequest(model: classificationModel, completionHandler: { [weak self] request, error in
                self?.processClassifications(for: request, error: error)
            })
            request.imageCropAndScaleOption = .scaleFit
            return request
        }
        catch {
            fatalError("Error! Can't use Model.")
        }
    }()
    func classifyImage(receivedImage: UIImage) {
        let orientation = CGImagePropertyOrientation(rawValue: UInt32(receivedImage.imageOrientation.rawValue))
        if let image = CIImage(image: receivedImage) {
            DispatchQueue.global(qos: .userInitiated).async {
                let handler = VNImageRequestHandler(ciImage: image, orientation: orientation!)
                do {
                    try handler.perform([self.classificationRequest])
                }
                catch {
                    fatalError("Error classifying image!")
                }
            }
        }
    }

The problem started when I tried to pass a UIImage to run inference on the model. The input type of the original model was MultiArray (Float32 1 x 224 x 224 x 3). Using Coremltools library I was able to convert the input type to Image (Color 224 x 224) using Python.

This worked and here is my code:

Code Block 
import coremltools
import coremltools.proto.FeatureTypes_pb2 as ft
spec = coremltools.utils.load_spec("model.mlmodel")
input = spec.description.input[0]
input.type.imageType.colorSpace = ft.ImageFeatureType.RGB
input.type.imageType.height = 224
input.type.imageType.width = 224
coremltools.utils.save_spec(spec, "newModel.mlmodel")

My problem now is with the output type. I want to be able to access the confidence of the classification as well as the result label of the classification. Again using coremltools I was able to to access the output description and I got this.

Code Block 
name: "scores"
type {
  multiArrayType {
    dataType: FLOAT32
  }
}

I am trying to change it this way:

Code Block 
f = open("dict.txt", "r")
labels = f.read()
class_labels = labels.splitlines()
print(class_labels)
class_labels = class_labels[1:] 
assert len(class_labels) == 57
for i, label in enumerate(class_labels):
  if isinstance(label, bytes):
    class_labels[i] = label.decode("utf8")
classifier_config = ct.ClassifierConfig(class_labels)
output = spec.description.output[0]
output.type = ft.DictionaryFeatureType

Unfortunately this is not working and I can't find information only that can help me... This I don't know what to do next.

Thank you for your help!