0

I'm trying to use Vision with a custom model I trained, but I don't see a way to get the bounding box where Vision detected it in the frame.

The model: I've trained the model using CreateML, and it can detect 2 specific items. I tested the model in CreateML with various images and it detects the 2 items correctly and places a box around them. So, shouldn't Vision be able to give me the bounding box as well?

func prepare() {
    do {
        let vnModel = try VNCoreMLModel(for: modelFile.model)
        let coreMlRequest = VNCoreMLRequest(model: vnModel,
                                            completionHandler: { (request, error) in
            guard
                let results = request.results
                    as? [VNClassificationObservation]  // is this the right cast?
            else { return }

            // how do I get the bounding box from the results?
        })
        vnRequests = [coreMlRequest]
    }
    catch {
        print(error)
    }
}

func run(arFrame: ARFrame) {
    do {
        let requestHandler = VNImageRequestHandler(cvPixelBuffer: arFrame.capturedImage,
                                                   options: [:])
        try requestHandler.perform(self.vnRequests)
    }
    catch {
        print(error)
    }
}
Andy Jazz
  • 49,178
  • 17
  • 136
  • 220
thedp
  • 8,350
  • 16
  • 53
  • 95

1 Answers1

2

Did you actually train an object detection model? Or a classification model?

You only get bounding boxes for the object detection model, not for a classifier.

Assuming you trained an object detector, the correct class is VNRecognizedObjectObservation, not VNClassificationObservation.

Matthijs Hollemans
  • 7,706
  • 2
  • 16
  • 23
  • I'm very new to CoreML, it's my first model. It's possible I didn't do it correctly. I labeled 50 images, by placing a rect around the items I would like to train. When I try the model inside CreateML, it finds the items and places an accurate rect around them, so I'm guessing it was trained for object detection. – thedp Sep 05 '20 at 11:14
  • 1
    That sounds correct. :-) You should be getting `VNRecognizedObjectObservation` results then. – Matthijs Hollemans Sep 06 '20 at 09:52
  • ‘VNRecognizedObjectObservation’ did the trick. Thank you. – thedp Sep 06 '20 at 14:33