ViewController.swift updated

ViewController.swift

import AVFoundation
import UIKit

class AVSampleBufferView: UIView {

    var bufferLayer: AVSampleBufferDisplayLayer {
        return layer as! AVSampleBufferDisplayLayer
    }

    override static var layerClass: AnyClass {
        return AVSampleBufferDisplayLayer.self
    }

}

class ViewController: UIViewController, AVCaptureVideoDataOutputSampleBufferDelegate {

    private let avManager = AVManager()

    private var coder: H264Coder?
    
    private let cameraView = AVSampleBufferView()

    private let encodedView = AVSampleBufferView()

    private let decodedView = AVSampleBufferView()

    private let stackView: UIStackView = {
        let view = UIStackView()
        view.translatesAutoresizingMaskIntoConstraints = false
        view.axis = .vertical
        view.distribution = .fillEqually
        return view
    }()

    override func viewDidLoad() {
        super.viewDidLoad()
        view.addSubview(stackView)
        [cameraView, encodedView, decodedView].forEach {
            stackView.addArrangedSubview($0)
            $0.bufferLayer.videoGravity = .resizeAspect
        }
        cameraView.layer.backgroundColor = UIColor.red.cgColor
        encodedView.layer.backgroundColor = UIColor.green.cgColor
        decodedView.layer.backgroundColor = UIColor.blue.cgColor

        NSLayoutConstraint.activate([
            stackView.leadingAnchor.constraint(equalTo: view.leadingAnchor),
            stackView.trailingAnchor.constraint(equalTo: view.trailingAnchor),
            stackView.topAnchor.constraint(equalTo: view.topAnchor),
            stackView.bottomAnchor.constraint(equalTo: view.bottomAnchor)
        ])
    }

    override func viewDidAppear(_ animated: Bool) {
        super.viewDidAppear(animated)
        avManager.start(delegate: self)
    }

    lazy var coderCallback: (CMSampleBuffer) -> Void = { [weak self] sb in
        guard let self = self else { return }
        DispatchQueue.main.async {
            self.encodedView.bufferLayer.enqueue(sb)
        }
    }


    // MARK: - AVCaptureVideoDataOutputSampleBufferDelegate

    func captureOutput(_ output: AVCaptureOutput,
                       didOutput sampleBuffer: CMSampleBuffer,
                       from connection: AVCaptureConnection) {
        DispatchQueue.main.async {
            self.cameraView.bufferLayer.enqueue(sampleBuffer)
        }
        if coder == nil,
           let formatDescription = CMSampleBufferGetFormatDescription(sampleBuffer) {
            let dimens = formatDescription.dimensions
            coder = H264Coder(width: dimens.width, height: dimens.height)
            coder?.onFrame = coderCallback
        }
        coder?.encode(sampleBuffer)
    }

}