WWDC20 Optical Flow Sample Code

Hello ,

Here are some steps that might help you. the first thing is how to load the CIkernel for your visualization . it is difficult to find the documentation on it.

Normaly you should convert the "Core Image Kernel Language" function to "Metal Shading Language" but you can still load the kernel manually with a String (take care this function is deprecated)

Sample to load Kernel :

import Foundation import CoreImage

class OpticalFlowVisualizerFilter: CIFilter {
   
  var inputImage: CIImage?
  let callback: CIKernelROICallback = {
      (index, rect) in
        return rect
      }
   
  static var kernel: CIKernel = { () -> CIKernel in
    
    /*let url = Bundle.main.url(forResource: "OpticalFlowVisualizer",
                 withExtension: "ci.metal")!
    let data = try! Data(contentsOf: url)
     
    return try! CIKernel(functionName: "flowView2",
                 fromMetalLibraryData: data)*/
    var source = "//\n// OpticalFlowVisualizer.cikernel\n// SampleVideoCompositionWithCIFilter\n//\n\n\nkernel vec4 flowView2(sampler image, float minLen, float maxLen, float size, float tipAngle)\n{\n\t/// Determine the color by calculating the angle from the .xy vector\n\t///\n\tvec4 s = sample(image, samplerCoord(image));\n\tvec2 vector = s.rg - 0.5;\n\tfloat len = length(vector);\n\tfloat H = atan(vector.y,vector.x);\n\t// convert hue to a RGB color\n\tH *= 3.0/3.1415926; // now range [3,3)\n\tfloat i = floor(H);\n\tfloat f = H-i;\n\tfloat a = f;\n\tfloat d = 1.0 - a;\n\tvec4 c;\n\t\t if (H<-3.0) c = vec4(0, 1, 1, 1);\n\telse if (H<-2.0) c = vec4(0, d, 1, 1);\n\telse if (H<-1.0) c = vec4(a, 0, 1, 1);\n\telse if (H<0.0) c = vec4(1, 0, d, 1);\n\telse if (H<1.0) c = vec4(1, a, 0, 1);\n\telse if (H<2.0) c = vec4(d, 1, 0, 1);\n\telse if (H<3.0) c = vec4(0, 1, a, 1);\n\telse       c = vec4(0, 1, 1, 1);\n\t// make the color darker if the .xy vector is shorter\n\tc.rgb *= clamp((len-minLen)/(maxLen-minLen), 0.0,1.0);\n\t/// Add arrow shapes based on the angle from the .xy vector\n\t///\n\tfloat tipAngleRadians = tipAngle * 3.1415/180.0;\n\tvec2 dc = destCoord(); // current coordinate\n\tvec2 dcm = floor((dc/size)+0.5)*size; // cell center coordinate\n\tvec2 delta = dcm - dc; // coordinate relative to center of cell\n\t// sample the .xy vector from the center of each cell\n\tvec4 sm = sample(image, samplerTransform(image, dcm));\n\tvector = sm.rg - 0.5;\n\tlen = length(vector);\n\tH = atan(vector.y,vector.x);\n\tfloat rotx, k, sideOffset, sideAngle;\n\t// these are the three sides of the arrow\n\trotx = delta.x*cos(H) - delta.y*sin(H);\n\tsideOffset = size*0.5*cos(tipAngleRadians);\n\tk = 1.0 - clamp(rotx-sideOffset, 0.0, 1.0);\n\tc.rgb *= k;\n\tsideAngle = (3.14159 - tipAngleRadians)/2.0;\n\tsideOffset = 0.5 * sin(tipAngleRadians / 2.0);\n\trotx = delta.x*cos(H-sideAngle) - delta.y*sin(H-sideAngle);\n\tk = clamp(rotx+size*sideOffset, 0.0, 1.0);\n\tc.rgb *= k;\n\trotx = delta.x*cos(H+sideAngle) - delta.y*sin(H+sideAngle);\n\tk = clamp(rotx+ size*sideOffset, 0.0, 1.0);\n\tc.rgb *= k;\n\t/// return the color premultiplied\n\tc *= s.a;\n\treturn c;\n}"
    return try! CIKernel(source: source)!
  }()

  override var outputImage : CIImage? {
    get {
      guard let input = inputImage else {return nil}
      return OpticalFlowVisualizerFilter.kernel.apply(extent: input.extent, roiCallback: callback, arguments: [input, 0.0, 100.0, 10.0, 30.0])
    }
  }
}

Then , The optical flow works with a pair of frames. You can use AVfoundation to extract frame from your video like this.

self.videoAssetReaderOutput = AVAssetReaderTrackOutput(track: self.videoTrack, outputSettings: [kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_420YpCbCr8BiPlanarFullRange])
    guard self.videoAssetReaderOutput != nil else {
      return false
    }

You can get CVPixelBuffer for your VNRequest like this

func nextFrame() -> CVPixelBuffer? {
    guard let sampleBuffer = self.videoAssetReaderOutput.copyNextSampleBuffer()
    else {
      return nil
    }
    currentFrame += 1
    return CMSampleBufferGetImageBuffer(sampleBuffer)
  }

And compare 2 frames like this :

     
    var requestHandler = VNSequenceRequestHandler()
    var previousImage = ofRequest.previousImage
    var observationImage: CIImage?
    
    let visionRequest = VNGenerateOpticalFlowRequest(targetedCIImage: ofRequest.targetImage, options: [:])
     
    do {
      try requestHandler.perform([visionRequest], on: previousImage)
      if let pixelBufferObservation = visionRequest.results?.first as? VNPixelBufferObservation
      {
        observationImage = CIImage(cvImageBuffer: pixelBufferObservation.pixelBuffer)
      }
    } catch {
      print(error)
    }
    
    let ciFilter = OpticalFlowVisualizerFilter()
    ciFilter.inputImage = observationImage
    let output = ciFilter.outputImage
    return output!
  }

I hope this helps you a bit