AVFrame to PNG

gistfile1.txt

func savePNG(_ frame: AVFrame, to url: String) throws {
    let fmtCtx = try AVFormatContext(format: nil, filename: url)

    guard let codec = AVCodec.findEncoderById(.PNG) else {
        fatalError("png codec doesn't exist")
    }
    guard let codecCtx = AVCodecContext(codec: codec) else {
        fatalError("Could not allocate video codec context")
    }
    codecCtx.width = frame.width
    codecCtx.height = frame.height
    codecCtx.pixFmt = .RGB24
    codecCtx.timebase = AVRational(num: 1, den: 1)
    codecCtx.framerate = AVRational(num: 0, den: 1)
    try codecCtx.openCodec()

    guard let stream = fmtCtx.addStream(codec: codec) else {
        fatalError("Failed allocating output stream")
    }
    try stream.copyParameters(from: codecCtx)

    if !fmtCtx.oformat!.flags.contains(.noFile) {
        try fmtCtx.openIO(url: url, flags: .write)
    }

    try fmtCtx.writeHeader()

    let dstPkt = AVPacket()
    let dstFrame = AVFrame()
    dstFrame.width = codecCtx.width
    dstFrame.height = codecCtx.height
    dstFrame.pixFmt = codecCtx.pixFmt

    let swsCtx = SwsContext(
        srcW: frame.width,
        srcH: frame.height,
        srcFormat: frame.pixFmt,
        dstW: dstFrame.width,
        dstH: dstFrame.height,
        dstFormat: dstFrame.pixFmt,
        flags: .bilinear
    )!
    let srcSlice = frame.data.map({ UnsafePointer($0) })
    let srcStride = frame.linesize.map({ Int32($0) })
    // buffer is going to be written to rawvideo file, no alignment
    let dstData = UnsafeMutablePointer<UnsafeMutablePointer<UInt8>?>.allocate(capacity: 4)
    let dstLinesize = UnsafeMutablePointer<Int32>.allocate(capacity: 4)
    defer {
        dstData.deallocate()
        dstLinesize.deallocate()
    }
    let dstBufSize = AVImage.alloc(
        data: dstData,
        linesizes: dstLinesize,
        width: dstFrame.width,
        height: dstFrame.height,
        pixFmt: dstFrame.pixFmt,
        align: 1
    )
    if dstBufSize < 0 {
        fatalError("Could not allocate destination image")
    }
    defer { AVImage.free(dstData) }

    swsCtx.scale(
        src: srcSlice,
        srcStride: srcStride,
        srcSliceY: 0,
        srcSliceH: frame.height,
        dst: dstData,
        dstStride: dstLinesize
    )

    dstFrame.data = [dstData[0], dstData[1], dstData[2], dstData[3]]
    dstFrame.linesize = [Int(dstLinesize[0]), Int(dstLinesize[1]), Int(dstLinesize[2]), Int(dstLinesize[3])]

    do {
        try codecCtx.sendFrame(dstFrame)
    } catch {
        fatalError("Error while sending a packet to the decoder: \(error)")
    }
    while true {
        do {
            try codecCtx.receivePacket(dstPkt)
        } catch let err as AVError where err == .EAGAIN || err == .EOF {
            break
        } catch {
            fatalError("Error while receiving a packet from the encoder: \(error)")
        }

        dstPkt.streamIndex = 0
        dstPkt.pts = .noPTS
        dstPkt.dts = .noPTS
        dstPkt.duration = 0
        dstPkt.pos = -1

        try fmtCtx.interleavedWriteFrame(pkt: dstPkt)

        dstPkt.unref()
    }
    dstFrame.unref()

    try fmtCtx.writeTrailer()
}