LoginSignup
8
6

More than 3 years have passed since last update.

@pic18f14k50

Vision+CoreML+Metal で信号機の点灯色を読み取ってみる

Posted at

色の表現としてHSV色空間というものに触れる機会があったので、信号機の色を読み取ってみた。
信号機の画像認識には YOLOv3-Tiny のCore MLモデルを利用。
＜完成イメージ＞
demo.png demo.gif
※このgif画像は外で撮影した動画をモニタで再生し、それをiPhoneで撮影。

「色」をどう読み取るのか　〜HSV色空間について〜

業務でアプリ開発をしていて色の指定にはRGBを使っていて、
たまに見かける↓での色表現(HSV色空間)は 何が嬉しいのか？？ と思っていた。

HSV色空間とは（Wikipediaより）

HSVモデル(英: HSV model)は色相(Hue)、彩度(Saturation・Chroma)、明度(Value・Brightness)の三つの成分からなる色空間。

　嬉しいことは色だけに色々あるようで、0~360度の色の輪(色相環)で表現される「色相（Hue）」を利用することで「赤っぽい色」「青っぽい色」といった人間の認識に近い処理ができる。

　ちなみに、AppleのCIImageのフィルターのドキュメント 『Applying a Chroma Key Effect』 ではクロマキー合成のため緑っぽい背景色を消すのに色相の108~144度で判定している。

　で、HSVを使って何かやるなら信号機読み取ってみるか、 ということで信号機が何色に点灯しているのかの判別をしてみた。以下、作成したサンプルプログラムについて解説します。

信号機の点灯色の識別

色相

　点灯色は次の色相で判定（家の近所の３つの信号機、かつ、天気の良い日中で判定できるように指定したので汎用性はないです）。

点灯色 色相 備考
青 140〜220度 「青」といっても「緑」に近い場合がある
黄 5〜60度 「黄」もかなり「赤」に踏み込まないと認識せず。
赤 340〜5度 比較的安定して認識できる
彩度

　色相だけだと信号機の枠の部分とかのランプ以外のところが 若干の青み があるだけで、青信号と認識してしまうので「彩度(Saturation)」を使って、しっかりとした青色 を判定できるようにしている。

彩度 備考
80%以上 【彩度とは】
色の鮮やかさの尺度（0~100%）。
白・グレー・黒（つまり、R = G = B）は鮮やかではないので 0%。
R, G, B すべて0ではないが、0があるのとき彩度が 100%

＜彩度による判定効果の実例＞
色相環（ここで撮影した色相環は中央に行くについれて白っぽくなる＝彩度が落ちる）をiPhoneで撮影して、青・黄・赤の各色相を単色でぬり潰す加工（それ以外の色は黒）を行い、彩度別に表示範囲を絞ったものがこちら。

彩度
指定なし		 彩度
30%以上を表示		 彩度
80%以上を表示
s_0.PNG s_30.PNG s_80.PNG

　彩度の指定がない場合、何かしらの色に分類されてしまい、この例では白い部分が黄色に判定されている。彩度の大きな部分だけ表示することで、より鮮やかな部分だけ残る、という効果が確認できる。

色の読み取りステップ

  1. カメラ画像をキャプチャする
  2. 画像からオブジェクト（信号機）を識別する
  3. 信号機部分の画像だけ切り取る
  4. 3.にブラーをかけて画像内の色のばらつきを抑える
  5. 4.の画像の青信号・黄信号・赤信号の色区別する
  6. 何色が点灯しているのか判定する

1. カメラ画像をキャプチャする

これはAppleの『YOLOv3-Tinyのサンプルプログラム』と同じ。一点、キャプチャした画像からあとで画像の切り出しがしやすいようにCMSampleBuffer から CGImage に変換しておく。

if let pb = CMSampleBufferGetImageBuffer(sampleBuffer) {
    // 検出した信号機をキャプチャー画像の上に重ねて表示させるためCGImageにしてとっておく
    let ciImage = CIImage(cvPixelBuffer: pb)
    let pixelBufferWidth = CGFloat(CVPixelBufferGetWidth(pixelBuffer))
    let pixelBufferHeight = CGFloat(CVPixelBufferGetHeight(pixelBuffer))
    let imageRect:CGRect = CGRect(x: 0,y: 0,width: pixelBufferWidth, height: pixelBufferHeight)
    let ciContext = CIContext.init()

    // Metal側でピクセルフォーマットを BGRA8 にしているので、CGImage作る時も合わせておく。
    cgimage = ciContext.createCGImage(ciImage, from: imageRect, format: .BGRA8, colorSpace: nil)
}

ポイントは、あとでMetalで加工する際のピクセルフォーマットを BGRA にしているので、CGImageを作る際にピクセルフォーマットを合わせているところ。

2. 画像からオブジェクト（信号機）を識別する

これも『YOLOv3-Tinyのサンプルプログラム』と同様。ここでは信号機以外の認識情報は捨てる。


let objectRecognition = VNCoreMLRequest(model: model, completionHandler: { (request, error) in
    DispatchQueue.main.async(execute: {
        self.cropedCGImage = nil

        guard let results = request.results else { return }
        for observation in results {
            guard let objectObservation = observation as? VNRecognizedObjectObservation,
                  objectObservation.labels[0].identifier == "traffic light" else {
                // 信号機以外の認識画像は捨てる
                continue
            }
// 略

3. 信号機部分の画像だけ切り取る

1.で取得しておいたキャプチャ画像から信号機部分の画像を切りだす。

// 信号機部分の画像だけ切り出す
let objectBounds = VNImageRectForNormalizedRect(objectObservation.boundingBox, self.videoWidth, self.videoHeight)
self.cropedCGImage = self.cgimage?.cropping(to: objectBounds)

あと、Metalで加工するため、MTLTexture に変換しておく。

let textureLoader = MTKTextureLoader(device: device)
// CGImageからMTLTexture作る時はusage指定が必要。
let usage = MTLTextureUsage.renderTarget.rawValue | MTLTextureUsage.shaderRead.rawValue | MTLTextureUsage.shaderWrite.rawValue
let texture = try? textureLoader.newTexture(cgImage: cropped,
      options: [.textureUsage : NSNumber.init(value:usage)])

ポイントはMTKTextureLoaderによる CGImage->MTLTextureへの変換時に、usage を設定するところ。これを設定しないと「usageを設定しろ」と実行時におこられる。shaderRead はシェーダー側で読めるようにする、renderTarget　は色の情報を取り扱えるようにする、という指定。shaderWrite は、このあとここで作成したMTLTextureに対し、Metal Performance Shadersでブラーをかけて書き込みをするので指定が必要。

4. 3.にブラーをかけて画像内の色のばらつきを抑える

このサンプルでの点灯色の判定方法は後述するが、画像上、いくつかの「点」の色を調べて、点灯色を判定している。LED信号機のように点が密集していたり、色のノイズがあると、判定精度が落ちるので、ブラーをかけて判定精度を上げる。

// LED信号機のつぶつぶを潰して信号判定精度を上げる
let blurKernek = MPSImageGaussianBlur(device: device, sigma: 2)
blurKernek.encode(commandBuffer: commandBuffer,
                  inPlaceTexture: &texture, fallbackCopyAllocator: nil)

ブラーの強さは適当。

5. 4.の画像の青信号・黄信号・赤信号の色区別する

ここはMetalのコンピュートシェーダーで処理。入力となるMTLTexture（信号機画像）から、色相・彩度により青信号は青(R,G,B=0,0,1)、黄信号は黄色(R,G,B=1,1,0)、赤信号は赤(R,G,B=1,0,0)に変換したMTLTextureを出力。

Shader.metal
// 色相・彩度に応じて、色を青・黄・青に塗りつぶす
kernel void colorClassify(texture2d<half, access::read> inTexture [[ texture(0) ]],
                          texture2d<half, access::write> outTexture [[ texture(1) ]],
                          uint2 gid [[thread_position_in_grid]]) {

    half3 color = inTexture.read(gid).rgb;

    half s = saturation(color.r, color.g, color.b);
    // 彩度が低い場合は判定NG
    if(!isGoodSaturation(s)) {
        outTexture.write(half4(0.0), gid);
        return;
    }

    half h = hue(color.r, color.g, color.b);
    if(isBlue(h)) {
        outTexture.write(half4(0.0, 0.0, 1.0, 1.0), gid);
        return;
    }
    if(isYellow(h)) {
        outTexture.write(half4(1.0, 1.0, 0.0, 1.0), gid);
        return;
    }
    if(isRed(h)) {
        outTexture.write(half4(1.0, 0.0, 0.0, 1.0), gid);
        return;
    }
    outTexture.write(half4(0.0), gid);
}

RGB→HSVへの変換は、こちらのサイト『RGBとHSVの相互変換[色見本/サンプル付き]』を参考にさせてもらいました。ちなみにwikiの「色相」を見るとこの変換方法は「大まかな値を求める」方法らしい。より正確な変換方法は、こちらのサイト『HSVの計算について』1 にある方法が良さそうだが、ちょっと面倒なのでスキップ。

ちなみに、
サンプルを作り始めたときは重い処理もあるだろうからと考えてMetalを使ったが、結局、識別した色を可視化するのに使っただけで、点灯色を把握するだけならMetalを使う必要なかった。可視化するにしても CIColorKernel がお手軽でよかったかも。

6. 何色が点灯しているのか判定する

ここが一番悩ましかった。手法は無数にあるような気がするが、この辺り詳しくないので、最低限、次を実現できる方法を考えた。

  • 処理が軽いこと。
  • 青空の影響を受けにくいこと。

で、実装したのは次のように信号機の縦中央・横に７箇所の色を取得して、色の多数決で点灯色を判定する、と言う方法（７箇所にしたのは適当）。
＜判定箇所のイメージ＞
imga.png

これなら青空の影響（青信号の誤判定）を避けやすく、かつ、判定処理も軽め。

// MTLTextureをUInt8の配列でアクセスできるようにする
let bytesPerPixel: Int = 4
let imageByteCount = changedTexture.width * changedTexture.height * bytesPerPixel
let bytesPerRow = changedTexture.width * bytesPerPixel
var src = [UInt8](repeating: 0, count: Int(imageByteCount))
let region = MTLRegionMake2D(0, 0, changedTexture.width, changedTexture.height)
classifiedTexture.getBytes(&src, bytesPerRow: bytesPerRow, from: region, mipmapLevel: 0)

// 画像の（見た目の）縦中央の画素を、右から左に７箇所等間隔で色を取得し、多数決で点灯色を判定
var signalCount: [(UIColor, Int)] = [(.blue, 0), (.yellow, 0), (.red, 0)]
let step = changedTexture.height / 8
for y in stride(from: step, to: changedTexture.height - step, by: step) {
    // 【注意】①BGRAの順番に色があること　②横向き画像であるといこと
    let baseIndex = (y * changedTexture.width + (changedTexture.width / 2)) * bytesPerPixel
    if src[baseIndex + 0] != 0 {
        signalCount[0].1 += 1   // 青＋１
    } else if src[baseIndex + 1] != 0 {
        signalCount[1].1 += 1   // 黄＋１ ※黄色＝赤＋緑なので赤より先に判定すること。
    } else if src[baseIndex + 2] != 0 {
        signalCount[2].1 += 1   // 赤＋１
    }
}

なお、MTLTextureから画素を取り出す方法については、こちらの記事『Metalシェーディング言語を使用したUIImageの加工』を参考にさせていただきました。

あと、上記の方法だけだと、信号機の画像認識ができなかったり、色のノイズで誤判定されたり、信号機を撮影するとフリッカーがでたり、で安定して認識できなかったので、過去１０回の個別の点灯色の認識結果の中で一番多く認識できた色＝点灯色、という判定も入れた。


if signalHistory.count > 10 {
    // 過去の判定結果は１０個まで
    signalHistory.removeFirst()
}

var signalCount: [(UIColor, Int)] = [(.blue, 0), (.yellow, 0), (.red, 0)]
signalHistory.forEach { signal in
    switch signal {
    case .blue:   signalCount[0].1 += 1
    case .yellow: signalCount[1].1 += 1
    case .red:    signalCount[2].1 += 1
    default: break
    }
}

if let signal = signalCount.max(by: { $0.1 < $1.1 }) {
    return signal.0
} else {
    return .clear
}

いろいろと条件が整っているときには点灯色は読み取れている。
（安定して認識できないケースは多々ある。最後に課題を記載）

最後に

目的がHSV色空間を使って何かする、、なので目的は達成したが、真剣に信号機の色を読み取るとなると課題は多い。思いつく課題は次の通り。

  1. そもそも信号機の画像認識精度がよくない（YOLOv3-Tinyは日本の信号機で学習されてないのでしょう）。日本の信号機を使って、信号機だけで学習すれば改善されそう。
  2. 色の判定を、ベタに色で判定した。これだと、夜とか、逆光とか撮影条件による影響が大きいように思われる。撮影条件×各色を含めた信号機の機械学習されたモデルがあれば、色で判定しなくていいかもしれない。画像判定前（＆トレーニングデータ画像）に適切な前処理は必要そう。
  3. 撮影すると信号機って結構小さい。信号機を識別させるにはより信号機が大きく見える条件付けが必要かも。例えば、画像の上半分にしか信号機は存在しない前提を置くとか。

全体ソースコード

ViewController.swift
import UIKit
import AVFoundation
import Vision
import Metal
import MetalKit
import MetalPerformanceShaders

class ViewController: UIViewController {

    @IBOutlet weak var preview: UIView!
    @IBOutlet weak var lblColor: UILabel!

    // キャプチャ画像サイズ
    private var videoWidth = 0
    private var videoHeight = 0
    // キャプチャ画像の入出力
    private var rootLayer: CALayer! = nil
    private let session = AVCaptureSession()
    private let videoDataOutput = AVCaptureVideoDataOutput()
    private let videoDataOutputQueue = DispatchQueue(label: "VideoDataOutput",
                                                     qos: .userInitiated,
                                                     attributes: [],
                                                     autoreleaseFrequency: .workItem)
    private var previewLayer: AVCaptureVideoPreviewLayer! = nil

    private var request: VNCoreMLRequest?
    // 信号機画像を加工した画像を乗せるレイヤー
    private var detectionOverlay: CALayer! = nil
    private var cgimage: CGImage?
    // Metal
    private let device = MTLCreateSystemDefaultDevice()!
    private lazy var commandQueue = device.makeCommandQueue()!
    // 描画用
    private var pipelineState: MTLRenderPipelineState!
    private var imagePlaneVertexBuffer: MTLBuffer!
    private var timer: CADisplayLink!
    private var metalLayer: CAMetalLayer!
    // 色分類用
    private var colorClassComputeState: MTLComputePipelineState!
    private var threadgroupSize = MTLSizeMake(16, 16, 1)
    private var classifiedTexture: MTLTexture!
    // Metal 画像頂点座標(x,y), uv座標(u,v)
    private let kImagePlaneVertexData: [Float] = [
        -1.0, -1.0, 0.0, 1.0,
        1.0, -1.0, 1.0, 1.0,
        -1.0, 1.0, 0.0, 0.0,
        1.0, 1.0, 1.0, 0.0
    ]
    // 信号機部分を切り出した画像
    var cropedCGImage: CGImage?
    // 点灯色の切り出し記録。チラつき抑止のために使う。
    var signalHistory: [UIColor] = []

    override func viewDidLoad() {
        super.viewDidLoad()

        setupAVCapture()
        setupLayers()
        setupVision()
        setupMetal()
        self.view.bringSubviewToFront(lblColor)
        // キャプチャ開始
        session.startRunning()
        // CAMetalLayerを使うので描画タイミングを生成
        timer = CADisplayLink(target: self, selector: #selector(ViewController.drawLoop))
        timer.add(to: RunLoop.main, forMode: .default)
    }

    override func viewDidLayoutSubviews() {
        super.viewDidLayoutSubviews()
        // 画像認識は 640x480 で行うのでLayer全体を右９０度回転する
        let bounds = rootLayer.bounds
        let xScale = bounds.size.width / CGFloat(videoHeight)
        let yScale = bounds.size.height / CGFloat(videoWidth)
        let scale = fmax(xScale, yScale)
        detectionOverlay.setAffineTransform(CGAffineTransform(rotationAngle: CGFloat(.pi / 2.0)).scaledBy(x: scale, y: scale))
        detectionOverlay.position = CGPoint(x: bounds.midX, y: bounds.midY)
    }

    // CAMetalLayerに信号機画像を描画
    @objc func drawLoop() {
        guard let cropped = cropedCGImage else {
            metalLayer.isHidden = true
            return
        }
        metalLayer.isHidden = false
        autoreleasepool {
            let textureLoader = MTKTextureLoader(device: device)
            // CGImageからMTLTexture作る時はusage指定が必要。
            let usage = MTLTextureUsage.renderTarget.rawValue | MTLTextureUsage.shaderRead.rawValue | MTLTextureUsage.shaderWrite.rawValue
            let texture = try? textureLoader.newTexture(cgImage: cropped,
                  options: [.textureUsage : NSNumber.init(value:usage)])

            // 切り抜いた信号機画像を点灯色が目立つように加工
            guard let changedTexture = changeColor(texture: texture) else {
                // 画像変換に失敗しているので点灯色識別結果に青・黄・赤以外を設定
                signalHistory.append(.clear)
                return
            }

            // MTLTextureをUInt8の配列でアクセスできるようにする
            let bytesPerPixel = 4
            let imageByteCount = changedTexture.width * changedTexture.height * bytesPerPixel
            let bytesPerRow = changedTexture.width * bytesPerPixel
            var src = [UInt8](repeating: 0, count: Int(imageByteCount))
            let region = MTLRegionMake2D(0, 0, changedTexture.width, changedTexture.height)
            classifiedTexture.getBytes(&src, bytesPerRow: bytesPerRow, from: region, mipmapLevel: 0)

            // 画像の（見た目の）縦中央の画素を、右から左に７箇所等間隔で色を取得し、多数決で点灯色を判定
            var signalCount: [(UIColor, Int)] = [(.blue, 0), (.yellow, 0), (.red, 0)]
            let step = changedTexture.height / 8
            for y in stride(from: step, to: changedTexture.height - step, by: step) {
                // 【注意】①BGRAの順番に色があること　②横向き画像であるといこと
                let baseIndex = (y * changedTexture.width + (changedTexture.width / 2)) * bytesPerPixel
                if src[baseIndex + 0] != 0 {
                    signalCount[0].1 += 1   // 青＋１
                } else if src[baseIndex + 1] != 0 {
                    signalCount[1].1 += 1   // 黄＋１ ※黄色＝赤＋緑なので赤より先に判定すること。
                } else if src[baseIndex + 2] != 0 {
                    signalCount[2].1 += 1   // 赤＋１
                }
            }

            if let signal = signalCount.max(by: { $0.1 < $1.1 }), signal.1 > 0 {
                signalHistory.append(signal.0)
            } else {
                signalHistory.append(.clear)
            }

            // 過去１０回の点灯色認識結果の中で一番多い点灯色を採用
            let detectedColor = detectedColorWithouitNoise()
            lblColor.text = detectedColor.jaText
            lblColor.backgroundColor = detectedColor
            lblColor.textColor = detectedColor.textColor

            // 加工済み信号機画像を描画
            drawSignal()
        }
    }

    private func detectedColorWithouitNoise() -> UIColor {
        if signalHistory.count > 10 {
            // 過去の判定結果は１０個まで
            signalHistory.removeFirst()
        }

        var signalCount: [(UIColor, Int)] = [(.blue, 0), (.yellow, 0), (.red, 0)]
        signalHistory.forEach { signal in
            switch signal {
            case .blue:   signalCount[0].1 += 1
            case .yellow: signalCount[1].1 += 1
            case .red:    signalCount[2].1 += 1
            default: break
            }
        }

        if let signal = signalCount.max(by: { $0.1 < $1.1 }) {
            return signal.0
        } else {
            return .clear
        }
    }

    private func changeColor(texture: MTLTexture?) -> MTLTexture? {
        guard var texture = texture else { return nil }
        let commandBuffer = self.commandQueue.makeCommandBuffer()!

        // LED信号機のつぶつぶを潰して信号判定精度を上げる
        let blurKernek = MPSImageGaussianBlur(device: device, sigma: 2)
        blurKernek.encode(commandBuffer: commandBuffer,
                          inPlaceTexture: &texture, fallbackCopyAllocator: nil)

        // 色相・彩度により信号機画像を青・黄・赤の単色に塗りつぶす
        let computeEncoder = commandBuffer.makeComputeCommandEncoder()!
        let colorDesc = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .bgra8Unorm,
                                                                 width: texture.width, height: texture.height, mipmapped: false)
        colorDesc.usage = [.shaderRead, .shaderWrite]
        classifiedTexture = device.makeTexture(descriptor: colorDesc)

        let threadCountW = (texture.width + self.threadgroupSize.width - 1) / self.threadgroupSize.width
        let threadCountH = (texture.height + self.threadgroupSize.height - 1) / self.threadgroupSize.height
        let threadgroupCount = MTLSizeMake(threadCountW, threadCountH, 1)

        computeEncoder.setComputePipelineState(colorClassComputeState)
        computeEncoder.setTexture(texture, index: 0)
        computeEncoder.setTexture(classifiedTexture, index: 1)
        computeEncoder.dispatchThreadgroups(threadgroupCount, threadsPerThreadgroup: threadgroupSize)
        computeEncoder.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        return classifiedTexture
    }

    private func drawSignal() {
        guard let drawable = metalLayer.nextDrawable() else { return  }

        let commandBuffer = self.commandQueue.makeCommandBuffer()!

        let renderPassDescriptor = MTLRenderPassDescriptor()
        renderPassDescriptor.colorAttachments[0].texture = drawable.texture
        renderPassDescriptor.colorAttachments[0].loadAction = .clear
        renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColorMake(0.0, 0.0, 0.0, 1.0)

        let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor)!
        renderEncoder.setCullMode(.none)
        renderEncoder.setRenderPipelineState(pipelineState)
        renderEncoder.setVertexBuffer(imagePlaneVertexBuffer, offset: 0, index: 0)
        renderEncoder.setFragmentTexture(classifiedTexture, index: 0)
        renderEncoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
        renderEncoder.endEncoding()

        commandBuffer.present(drawable)
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
    }
}

extension ViewController {

    private func setupAVCapture() {
        let videoDevice = AVCaptureDevice.DiscoverySession(deviceTypes: [.builtInWideAngleCamera],
                                                           mediaType: .video,
                                                           position: .back).devices.first
        guard let deviceInput = try? AVCaptureDeviceInput(device: videoDevice!) else { return }
        // capture セッション セットアップ
        session.beginConfiguration()
        session.sessionPreset = .vga640x480

        // 入力デバイス指定
        session.addInput(deviceInput)

        // 出力先の設定
        session.addOutput(videoDataOutput)
        videoDataOutput.alwaysDiscardsLateVideoFrames = true
        videoDataOutput.videoSettings = [kCVPixelBufferPixelFormatTypeKey as String: Int(kCVPixelFormatType_420YpCbCr8BiPlanarFullRange)]
        videoDataOutput.setSampleBufferDelegate(self, queue: videoDataOutputQueue)
        let captureConnection = videoDataOutput.connection(with: .video)
        captureConnection?.isEnabled = true

        // ビデオの画像サイズ取得
        try? videoDevice!.lockForConfiguration()
        let dimensions = CMVideoFormatDescriptionGetDimensions((videoDevice?.activeFormat.formatDescription)!)
        videoWidth = Int(dimensions.width)
        videoHeight = Int(dimensions.height)
        videoDevice!.unlockForConfiguration()

        session.commitConfiguration()

        // プレビューセットアップ
        previewLayer = AVCaptureVideoPreviewLayer(session: session)
        previewLayer.videoGravity = AVLayerVideoGravity.resizeAspectFill
        rootLayer = preview.layer
        previewLayer.frame = rootLayer.bounds
        rootLayer.addSublayer(previewLayer)
    }

    private func setupVision() {
        guard let model = try? VNCoreMLModel(for: YOLOv3Tiny(configuration: MLModelConfiguration()).model) else { return }

        let objectRecognition = VNCoreMLRequest(model: model, completionHandler: { (request, error) in
            DispatchQueue.main.async(execute: {
                self.cropedCGImage = nil

                guard let results = request.results else { return }
                for observation in results {
                    guard let objectObservation = observation as? VNRecognizedObjectObservation,
                          objectObservation.labels[0].identifier == "traffic light" else {
                        // 信号機以外の認識画像は捨てる
                        continue
                    }

                    // 信号機部分の画像だけ切り出す
                    let objectBounds = VNImageRectForNormalizedRect(objectObservation.boundingBox, self.videoWidth, self.videoHeight)
                    self.cropedCGImage = self.cgimage?.cropping(to: objectBounds)
                    // 信号機画像（加工済み）を表示する位置を設定する
                    self.metalLayer.bounds = objectBounds
                    self.metalLayer.position = CGPoint(x: objectBounds.midX, y: objectBounds.midY)
                    break
                }
            })
        })
        request = objectRecognition
    }

    private func setupLayers() {
        // CAMetalLayerを乗せるレイヤー。あとで全体を９０度回転して使う。
        detectionOverlay = CALayer()
        detectionOverlay.bounds = CGRect(x: 0.0,
                                         y: 0.0,
                                         width: CGFloat(videoWidth),
                                         height: CGFloat(videoHeight))
        detectionOverlay.position = CGPoint(x: rootLayer.bounds.midX, y: rootLayer.bounds.midY)
        rootLayer.addSublayer(detectionOverlay)
        // Metalで作ったMTLTextureの描画レイヤー
        metalLayer = CAMetalLayer()
        metalLayer.device = device
        metalLayer.pixelFormat = .bgra8Unorm
        detectionOverlay.addSublayer(metalLayer)
    }

    private func setupMetal() {
        // 頂点座標バッファ確保＆頂点情報流し込み
        let vertexDataCount = kImagePlaneVertexData.count * MemoryLayout<Float>.size
        imagePlaneVertexBuffer = device.makeBuffer(bytes: kImagePlaneVertexData, length: vertexDataCount, options: [])
        // シェーダー指定
        let defaultLibrary = device.makeDefaultLibrary()!
        let pipelineStateDescriptor = MTLRenderPipelineDescriptor()
        pipelineStateDescriptor.vertexFunction = defaultLibrary.makeFunction(name: "imageVertex")!
        pipelineStateDescriptor.fragmentFunction = defaultLibrary.makeFunction(name: "imageFragment")!
        pipelineStateDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
        pipelineStateDescriptor.sampleCount = 1

        try? pipelineState = device.makeRenderPipelineState(descriptor: pipelineStateDescriptor)

        // 画像分類色設定
        let labelConvertShader = defaultLibrary.makeFunction(name: "colorClassify")!
        colorClassComputeState = try! self.device.makeComputePipelineState(function: labelConvertShader)
    }
}

extension ViewController: AVCaptureVideoDataOutputSampleBufferDelegate {
    func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
        // 信号機画像の検出開始
        guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { return }
        let imageRequestHandler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer,
                                                        orientation: .up,
                                                        options: [:])
        guard let request = self.request else { return }
        try? imageRequestHandler.perform([request])

        // 検出した信号機をキャプチャー画像の上に重ねて表示させるためCGImageにしてとっておく
        guard let pb = CMSampleBufferGetImageBuffer(sampleBuffer) else { return }
        let ciImage = CIImage(cvPixelBuffer: pb)
        let pixelBufferWidth = CGFloat(CVPixelBufferGetWidth(pixelBuffer))
        let pixelBufferHeight = CGFloat(CVPixelBufferGetHeight(pixelBuffer))
        let imageRect:CGRect = CGRect(x: 0,y: 0,width: pixelBufferWidth, height: pixelBufferHeight)
        let ciContext = CIContext.init()

        // Metal側でピクセルフォーマットを BGRA8 にしているので、CGImage作る時も合わせておく。
        cgimage = ciContext.createCGImage(ciImage, from: imageRect, format: .BGRA8, colorSpace: nil)
    }
}

extension UIColor {
    var jaText: String {
        switch self {
        case .blue: return "青"
        case .yellow: return "黄"
        case .red: return "赤"
        default: return ""
        }
    }

    var textColor: UIColor {
        switch self {
        case .blue: return .white
        case .yellow: return .black
        case .red: return .white
        default: return .black
        }
    }
}
Shader.metal
#include <metal_stdlib>
using namespace metal;

typedef struct {
    float2 position;
    float2 texCoord;
} ImageVertex;

typedef struct {
    float4 position [[position]];
    float2 texCoord;
} ColorInOut;

vertex ColorInOut imageVertex(const device ImageVertex* vertices [[ buffer(0) ]],
                              unsigned int vid [[ vertex_id ]]) {
    ColorInOut out;
    const device ImageVertex& cv = vertices[vid];
    out.position = float4(cv.position, 0.0, 1.0);
    out.texCoord = cv.texCoord;
    return out;
}

// rgb -> 色相変換
half hue(half r, half g, half b) {
    half rgbMax = max3(r, g, b);
    half rgbMin = min3(r, g, b);
    if(rgbMax == rgbMin) {
        return 0;
    } else {
        half hue;
        if(rgbMax == r) {
            hue = 60.0 * (g - b) / (rgbMax - rgbMin);
        } else if(rgbMax == g) {
            hue = 60.0 * (b - r) / (rgbMax - rgbMin) + 120;
        } else {
            hue = 60.0 * (r - g) / (rgbMax - rgbMin) + 240;
        }
        if(hue < 0) {
            hue += 360;
        }
        return hue;
    }
}

// rgb -> 彩度変換
half saturation(half r, half g, half b) {
    half rgbMax = max3(r, g, b);
    half rgbMin = min3(r, g, b);

    if(rgbMax == rgbMin) {
        return 0;
    } else {
        return (rgbMax - rgbMin) / rgbMax;
    }
}

bool isBlue(half hue) {
    return hue > 140 && hue < 220;
}

bool isYellow(half hue) {
    return hue > 5 && hue < 60;
}

bool isRed(half hue) {
    return hue < 5 || hue > 340;
}

bool isGoodSaturation(half saturation) {
    return saturation > 0.8;
}

fragment half4 imageFragment(ColorInOut in [[ stage_in ]],
                             texture2d<half, access::sample> signalTexture [[ texture(0) ]]) {
    constexpr sampler colorSampler(mip_filter::linear,
                                   mag_filter::linear,
                                   min_filter::linear);

    half4 color = signalTexture.sample(colorSampler, in.texCoord.xy);
    return color;
}

// 色相・彩度に応じて、色を青・黄・青に塗りつぶす
kernel void colorClassify(texture2d<half, access::read> inTexture [[ texture(0) ]],
                          texture2d<half, access::write> outTexture [[ texture(1) ]],
                          uint2 gid [[thread_position_in_grid]]) {

    half3 color = inTexture.read(gid).rgb;

    half s = saturation(color.r, color.g, color.b);
    // 彩度が低い場合は判定NG
    if(!isGoodSaturation(s)) {
        outTexture.write(half4(0.0), gid);
        return;
    }

    half h = hue(color.r, color.g, color.b);
    if(isBlue(h)) {
        outTexture.write(half4(0.0, 0.0, 1.0, 1.0), gid);
        return;
    }
    if(isYellow(h)) {
        outTexture.write(half4(1.0, 1.0, 0.0, 1.0), gid);
        return;
    }
    if(isRed(h)) {
        outTexture.write(half4(1.0, 0.0, 0.0, 1.0), gid);
        return;
    }
    outTexture.write(half4(0.0), gid);
}

参考URL

  1. 最初、wikiにあるatan()使う方法で作り始めたが正しく変換ができないのに気がついて、このサイトに辿り着いた ↩

8
6
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
8
6