How to know what class the image belongs to?

Question

I am trying to play around with Vision and CoreML.

Code:

extension CaptureImageView {
    private func loadImage() {
        guard let inputImage = image else { return }
        predictImage = inputImage
        performImageClassification()
    }
    
    private func performImageClassification() {
        let image = self.predictImage
        let resizedImage = image?.resizeTo(size: CGSize(width: 224, height: 224))
        let buffer = resizedImage?.toBuffer()
        let output = try? model.prediction(image: buffer!)
        if let output = output {
            breed = output.classLabel
        }
    }
}

I would like to:

1. See if this is a dog in the photo
2. If it's a dog, show the breed
3. If it's not a dog, show some alert.

The problem: That this recognition at the moment recognizes everything :) But I would like to recognize only dogs.

So as mentioned, I am using Vision, CoreML, MobileNetV2.mlmodel from Apple.

Answer 1

This model detects 1000 classes of objects. About 200 of these are dog breeds. Figure out which classes are dogs and only show something if the detected class is one of the dog classes.

If you want to detect other types of dogs, you will have to train your own model.