module Codec.FFmpeg.Encode where
import Codec.FFmpeg.Common
import Codec.FFmpeg.Enums
import Codec.FFmpeg.Internal.Linear
import Codec.FFmpeg.Scaler
import Codec.FFmpeg.Types
import Codec.Picture
import Control.Applicative
import Control.Monad (when, void)
import Control.Monad.Error.Class
import Data.Bits
import Data.Maybe (fromMaybe)
import Data.Ord (comparing)
import Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as VM
import Foreign.C.String
import Foreign.C.Types
import Foreign.ForeignPtr (touchForeignPtr)
import Foreign.Marshal.Alloc
import Foreign.Marshal.Utils
import Foreign.Ptr
import Foreign.Storable
foreign import ccall "avcodec_find_encoder"
avcodec_find_encoder :: AVCodecID -> IO AVCodec
foreign import ccall "avcodec_find_encoder_by_name"
avcodec_find_encoder_by_name :: CString -> IO AVCodec
foreign import ccall "av_opt_set"
av_opt_set :: Ptr () -> CString -> CString -> CInt -> IO CInt
foreign import ccall "avcodec_encode_video2"
avcodec_encode_video2 :: AVCodecContext -> AVPacket -> AVFrame -> Ptr CInt
-> IO CInt
foreign import ccall "av_image_alloc"
av_image_alloc :: Ptr (Ptr CUChar) -> Ptr CInt -> CInt -> CInt
-> AVPixelFormat -> CInt -> IO CInt
foreign import ccall "av_freep"
av_freep :: Ptr (Ptr a) -> IO ()
foreign import ccall "av_guess_format"
av_guess_format :: CString -> CString -> CString -> IO AVOutputFormat
foreign import ccall "avformat_alloc_output_context2"
avformat_alloc_output_context :: Ptr AVFormatContext -> AVOutputFormat
-> CString -> CString -> IO CInt
foreign import ccall "avformat_new_stream"
avformat_new_stream :: AVFormatContext -> AVCodec -> IO AVStream
foreign import ccall "av_write_frame"
av_write_frame :: AVFormatContext -> AVPacket -> IO CInt
foreign import ccall "av_interleaved_write_frame"
av_interleaved_write_frame :: AVFormatContext -> AVPacket -> IO CInt
foreign import ccall "avformat_write_header"
avformat_write_header :: AVFormatContext -> Ptr AVDictionary -> IO CInt
foreign import ccall "av_write_trailer"
av_write_trailer :: AVFormatContext -> IO CInt
foreign import ccall "avio_open"
avio_open :: Ptr AVIOContext -> CString -> AVIOFlag -> IO CInt
foreign import ccall "avio_close"
avio_close :: AVIOContext -> IO CInt
foreign import ccall "avformat_free_context"
avformat_free_context :: AVFormatContext -> IO ()
foreign import ccall "av_image_fill_arrays"
av_image_fill_arrays :: Ptr (Ptr CUChar) -> Ptr CInt -> Ptr CUChar
-> AVPixelFormat -> CInt -> CInt -> CInt -> IO CInt
foreign import ccall "av_image_fill_linesizes"
av_image_fill_linesizes :: Ptr CInt -> AVPixelFormat -> CInt -> IO CInt
data EncodingParams =
EncodingParams { epWidth :: CInt
, epHeight :: CInt
, epFps :: Int
, epCodec :: Maybe AVCodecID
, epPixelFormat :: Maybe AVPixelFormat
, epPreset :: String
}
defaultH264 :: CInt -> CInt -> EncodingParams
defaultH264 w h = EncodingParams w h 30 (Just avCodecIdH264) Nothing "medium"
defaultParams :: CInt -> CInt -> EncodingParams
defaultParams w h = EncodingParams w h 30 Nothing Nothing ""
checkFlag :: Bits a => a -> a -> Bool
checkFlag flg = \x -> (flg .&. x) /= allZeroBits
where allZeroBits = clearBit (bit 0) 0
initStream :: EncodingParams -> AVFormatContext -> IO (AVStream, AVCodecContext)
initStream ep _
| (epWidth ep `rem` 2, epHeight ep `rem` 2) /= (0,0) =
throwError $ strMsg "Video dimensions must be multiples of two"
initStream ep oc = do
codec <- maybe (getOutputFormat oc >>= getVideoCodecID) return (epCodec ep)
cod <- avcodec_find_encoder codec
when (getPtr cod == nullPtr)
(errMsg "Couldn't find encoder")
st <- avformat_new_stream oc cod
getNumStreams oc >>= setId st . subtract 1
ctx <- getCodecContext st
setWidth ctx (epWidth ep)
setHeight ctx (epHeight ep)
let framePeriod = AVRational 1 (fromIntegral $ epFps ep)
setTimeBase ctx framePeriod
setPixelFormat ctx $ case epPixelFormat ep of
Just fmt -> fmt
Nothing
| codec == avCodecIdRawvideo -> avPixFmtRgb24
| codec == avCodecIdGif -> avPixFmtPal8
| otherwise -> avPixFmtYuv420p
needsHeader <- checkFlag avfmtGlobalheader <$>
(getOutputFormat oc >>= getFormatFlags)
when needsHeader $
getCodecFlags ctx >>= setCodecFlags ctx . (.|. codecFlagGlobalHeader)
when (not . null $ epPreset ep) . void $
withCString "preset" $ \kStr ->
withCString (epPreset ep) $ \vStr ->
getPrivData ctx >>= \pd -> av_opt_set pd kStr vStr 0
rOpen <- open_codec ctx cod nullPtr
when (rOpen < 0) (throwError $ strMsg "Couldn't open codec")
return (st, ctx)
initTempFrame :: EncodingParams -> AVPixelFormat -> IO AVFrame
initTempFrame ep fmt = do
frame <- frame_alloc_check
setPixelFormat frame fmt
setWidth frame (epWidth ep)
setHeight frame (epHeight ep)
setPts frame 0
if fmt == avPixFmtRgb8 || fmt == avPixFmtPal8
then do r <- av_image_fill_linesizes (hasLineSize frame) fmt (epWidth ep)
when (r < 0) (errMsg "Error filling temporary frame line sizes")
else frame_get_buffer_check frame 32
return frame
allocOutputContext :: FilePath -> IO AVFormatContext
allocOutputContext fname = do
oc <- alloca $ \ocTmp -> do
r <- withCString fname $ \fname' ->
avformat_alloc_output_context
ocTmp (AVOutputFormat nullPtr)
nullPtr fname'
when (r < 0)
(errMsg "Couldn't allocate output format context")
peek ocTmp
when (getPtr oc == nullPtr)
(errMsg "Couldn't allocate output AVFormatContext")
return oc
avio_open_check :: AVFormatContext -> String -> IO ()
avio_open_check oc fname =
do r <- withCString fname $ \cstr ->
avio_open (hasIOContext oc) cstr avioFlagWrite
when (r < 0) (errMsg "Error opening IO for writing")
avio_close_check :: AVFormatContext -> IO ()
avio_close_check oc = do r <- getIOContext oc >>= avio_close
when (r /= 0) (errMsg "Error closing IO")
encode_video_check :: AVCodecContext -> AVPacket -> Maybe AVFrame -> IO Bool
encode_video_check ctx pkt frame =
alloca $ \gotOutput -> do
r <- avcodec_encode_video2 ctx pkt frame' gotOutput
when (r < 0) (errMsg "Error encoding frame")
(> 0) <$> peek gotOutput
where frame' = fromMaybe (AVFrame nullPtr) frame
write_header_check :: AVFormatContext -> IO ()
write_header_check oc = do r <- avformat_write_header oc nullPtr
when (r < 0) (errMsg "Error writing header")
write_frame_check :: AVFormatContext -> AVPacket -> IO ()
write_frame_check oc pkt = do r <- av_write_frame oc pkt
when (r < 0) (errMsg "Error writing frame")
write_trailer_check :: AVFormatContext -> IO ()
write_trailer_check oc = do r <- av_write_trailer oc
when (r /= 0) (errMsg "Error writing trailer")
palettizeRGB8 :: EncodingParams -> V.Vector CUChar -> V.Vector CUChar
palettizeRGB8 ep = \pix -> V.create $
do let pix' = V.unsafeCast pix :: V.Vector (V3 CUChar)
m <- VM.new (numPix + 1024)
V.mapM_ (\i -> let p = searchPal $ fromIntegral <$> (pix' V.! i)
in VM.unsafeWrite m i p)
(V.enumFromN 0 numPix)
VM.set (VM.unsafeSlice numPix 1024 m) 0
return m
where numPix = fromIntegral $ epWidth ep * epHeight ep
pal :: V.Vector (V3 CInt)
pal = V.generate 256 $ \i' ->
let i = fromIntegral i'
in V3 ((i `shiftR` 5) * 36)
(((i `shiftR` 2) .&. 7) * 36)
((i .&. 3) * 85)
searchPal = fromIntegral . flip V.minIndexBy pal . comparing . qd
palettizeJuicy :: EncodingParams -> V.Vector CUChar -> V.Vector CUChar
palettizeJuicy ep pix =
let (pix', pal) = palettize (PaletteOptions MedianMeanCut doDither 256)
(mkImage $ V.unsafeCast pix)
pal' = V.map (\(V3 r g b) -> V4 b g r (255::CUChar))
(V.unsafeCast $ imageData pal)
in V.unsafeCast (imageData pix') V.++ V.unsafeCast pal'
where mkImage = Image (fromIntegral $ epWidth ep) (fromIntegral $ epHeight ep)
doDither = epPreset ep == "dither"
frameWriter :: EncodingParams -> FilePath -> IO (Maybe (Vector CUChar) -> IO ())
frameWriter ep fname = do
oc <- allocOutputContext fname
(st,ctx) <- initStream ep oc
dstFmt <- getPixelFormat ctx
dstFrame <- initTempFrame ep dstFmt
sws <- if dstFmt /= avPixFmtPal8 && dstFmt /= avPixFmtRgb8
then Just <$>
swsInit (ImageInfo (epWidth ep) (epHeight ep) avPixFmtRgb24)
(ImageInfo (epWidth ep) (epHeight ep) dstFmt)
swsBilinear
else return Nothing
pkt <- AVPacket <$> av_malloc (fromIntegral packetSize)
stIndex <- getStreamIndex st
avio_open_check oc fname
write_header_check oc
tb <- getTimeBase st
codecTB <- getCodecContext st >>= getTimeBase
isRaw <- checkFlag avfmtRawpicture <$> (getOutputFormat oc >>= getFormatFlags)
let palettizer | dstFmt == avPixFmtPal8 = Just $ palettizeJuicy ep
| dstFmt == avPixFmtRgb8 = Just $ palettizeRGB8 ep
| otherwise = Nothing
frameTime = av_rescale_q 1 codecTB tb
mkImage :: Vector CUChar -> Image PixelRGB8
mkImage = let [w,h] = map fromIntegral [epWidth ep, epHeight ep]
in Image w h . V.unsafeCast
resetPacket = do init_packet pkt
setData pkt nullPtr
setSize pkt 0
writePacket = do setStreamIndex pkt stIndex
write_frame_check oc pkt
copyDstData = copyDstDataAux
copyDstDataAux pixels =
void . V.unsafeWith pixels $ \ptr ->
av_image_fill_arrays (castPtr $ hasData dstFrame)
(hasLineSize dstFrame)
(castPtr ptr)
dstFmt
(epWidth ep)
(epHeight ep)
1
scaleToDst sws' pixels = void $ swsScale sws' (mkImage pixels) dstFrame
fillDst = maybe copyDstData scaleToDst sws
addRaw Nothing = return ()
addRaw (Just pixels) =
do resetPacket
getPacketFlags pkt >>= setPacketFlags pkt . (.|. avPktFlagKey)
setSize pkt (fromIntegral pictureSize)
getPts dstFrame >>= setPts dstFrame . (+ frameTime)
getPts dstFrame >>= setPts pkt
getPts dstFrame >>= setDts pkt
V.unsafeWith pixels $ \ptr -> do
setData pkt (castPtr ptr)
writePacket
addEncoded Nothing = do resetPacket
encode_video_check ctx pkt Nothing >>=
flip when (writePacket >> addEncoded Nothing)
addEncoded (Just pixels) =
do resetPacket
let pixels' = maybe pixels ($ V.unsafeCast pixels) palettizer
fillDst pixels'
getPts dstFrame >>= setPts dstFrame . (+ frameTime)
encode_video_check ctx pkt (Just dstFrame) >>= flip when writePacket
let (fp,_,_) = V.unsafeToForeignPtr pixels'
touchForeignPtr fp
addFrame = if isRaw then addRaw else addEncoded
go Nothing = do addFrame Nothing
write_trailer_check oc
_ <- codec_close ctx
with dstFrame av_frame_free
av_free (getPtr pkt)
avio_close_check oc
avformat_free_context oc
go (Just pixels) = addFrame (Just pixels)
return go