CreateML results are awful at app while it's perfect in playground

238 Views Asked by At

I am trying to train MLModel with image classification. I created an app to create images to use as training data(at the end the same process will be used to get predictions). I get CVPixelBuffer from AvCaptureSession, convert it to UIImage and save it to documents directory as JPEG. Later I label them and train MLModel with CreateML in the playground. Results are %100 in the playground since I have collected thousands of images.

But when i integrate this model in my app and feed it with the same way, results are awful. I get CVPixelBuffer, convert it UIImage(to crop) and convert the cropped image to CVPixelBuffer and give it to model. I have to convert UIImage to CVPixelBuffer because CoreML model only excepts CVPixelBuffer. I convert UIImage to CVPixelBuffer with this method:

func pixelBuffer(width: Int, height: Int) -> CVPixelBuffer? {
    var maybePixelBuffer: CVPixelBuffer?
    let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
                 kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue]
    let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                     width,
                                     height,
                                     kCVPixelFormatType_32ARGB,
                                     attrs as CFDictionary,
                                     &maybePixelBuffer)

    guard status == kCVReturnSuccess, let pixelBuffer = maybePixelBuffer else {
        return nil
    }

    CVPixelBufferLockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
    let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer)

    guard let context = CGContext(data: pixelData,
                                  width: width,
                                  height: height,
                                  bitsPerComponent: 8,
                                  bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer),
                                  space: CGColorSpaceCreateDeviceRGB(),
                                  bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
        else {
            return nil
    }

    UIGraphicsPushContext(context)
    context.translateBy(x: 0, y: CGFloat(height))
    context.scaleBy(x: 1, y: -1)
    self.draw(in: CGRect(x: 0, y: 0, width: width, height: height))
    UIGraphicsPopContext()

    CVPixelBufferUnlockBaseAddress(pixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
    return pixelBuffer
}

I think I am getting poor results because CoreML model doesn't like the converted CVPixelBuffer.

Does anyone have any suggestion?

1

There are 1 best solutions below

12
On

You don't need any of this stuff. Let's have a look at the docs:

class VNCoreMLRequest : VNImageBasedRequest

Firstly, VNImageBasedRequest contains field regionOfInterest: CGRect { get set } where rectangle is normalized and lower-left relative. So you do not need to crop! Simply specify ROI.

Secondly, VNCoreMLRequest itself has field var imageCropAndScaleOption: VNImageCropAndScaleOption { get set } where you can specify how to act when height/width ratio does not match expected one (center crop, scale to fit/fill).