use crate::{ codec::{EncoderApi, EncoderCfg}, hw, HW_STRIDE_ALIGN, }; use hbb_common::{ anyhow::{anyhow, Context}, config::HwCodecConfig, lazy_static, log, message_proto::{EncodedVideoFrame, EncodedVideoFrames, Message, VideoFrame}, ResultType, }; use hwcodec::{ decode::{DecodeContext, DecodeFrame, Decoder}, encode::{EncodeContext, EncodeFrame, Encoder}, ffmpeg::{CodecInfo, CodecInfos, DataFormat}, AVPixelFormat, Quality::{self, *}, RateContorl::{self, *}, }; use std::sync::{Arc, Mutex}; lazy_static::lazy_static! { static ref HW_ENCODER_NAME: Arc>> = Default::default(); } const CFG_KEY_ENCODER: &str = "bestHwEncoders"; const CFG_KEY_DECODER: &str = "bestHwDecoders"; const DEFAULT_PIXFMT: AVPixelFormat = AVPixelFormat::AV_PIX_FMT_YUV420P; const DEFAULT_TIME_BASE: [i32; 2] = [1, 30]; const DEFAULT_GOP: i32 = 60; const DEFAULT_HW_QUALITY: Quality = Quality_Default; const DEFAULT_RC: RateContorl = RC_DEFAULT; pub struct HwEncoder { encoder: Encoder, yuv: Vec, pub format: DataFormat, pub pixfmt: AVPixelFormat, } impl EncoderApi for HwEncoder { fn new(cfg: EncoderCfg) -> ResultType where Self: Sized, { match cfg { EncoderCfg::HW(config) => { let ctx = EncodeContext { name: config.codec_name.clone(), width: config.width as _, height: config.height as _, pixfmt: DEFAULT_PIXFMT, align: HW_STRIDE_ALIGN as _, bitrate: config.bitrate * 1000, timebase: DEFAULT_TIME_BASE, gop: DEFAULT_GOP, quality: DEFAULT_HW_QUALITY, rc: DEFAULT_RC, }; let format = match Encoder::format_from_name(config.codec_name.clone()) { Ok(format) => format, Err(_) => { return Err(anyhow!(format!( "failed to get format from name:{}", config.codec_name ))) } }; match Encoder::new(ctx.clone()) { Ok(encoder) => Ok(HwEncoder { encoder, yuv: vec![], format, pixfmt: ctx.pixfmt, }), Err(_) => Err(anyhow!(format!("Failed to create encoder"))), } } _ => Err(anyhow!("encoder type mismatch")), } } fn encode_to_message( &mut self, frame: &[u8], _ms: i64, ) -> ResultType { let mut msg_out = Message::new(); let mut vf = VideoFrame::new(); let mut frames = Vec::new(); for frame in self.encode(frame).with_context(|| "Failed to encode")? { frames.push(EncodedVideoFrame { data: frame.data, pts: frame.pts as _, ..Default::default() }); } if frames.len() > 0 { let frames = EncodedVideoFrames { frames: frames.into(), ..Default::default() }; match self.format { DataFormat::H264 => vf.set_h264s(frames), DataFormat::H265 => vf.set_h265s(frames), } msg_out.set_video_frame(vf); Ok(msg_out) } else { Err(anyhow!("no valid frame")) } } fn use_yuv(&self) -> bool { false } fn set_bitrate(&mut self, bitrate: u32) -> ResultType<()> { self.encoder.set_bitrate((bitrate * 1000) as _).ok(); Ok(()) } } impl HwEncoder { /// Get best encoders. /// /// # Parameter /// `force_reset`: force to refresh config. /// `write`: write to config file. /// /// # Return /// `CodecInfos`: infos. /// `bool`: whether the config is refreshed. pub fn best(force_reset: bool, write: bool) -> (CodecInfos, bool) { let config = get_config(CFG_KEY_ENCODER); if !force_reset && config.is_ok() { (config.unwrap(), false) } else { let ctx = EncodeContext { name: String::from(""), width: 1920, height: 1080, pixfmt: DEFAULT_PIXFMT, align: HW_STRIDE_ALIGN as _, bitrate: 0, timebase: DEFAULT_TIME_BASE, gop: DEFAULT_GOP, quality: DEFAULT_HW_QUALITY, rc: DEFAULT_RC, }; let encoders = CodecInfo::score(Encoder::avaliable_encoders(ctx)); if write { set_config(CFG_KEY_ENCODER, &encoders) .map_err(|e| log::error!("{:?}", e)) .ok(); } (encoders, true) } } pub fn current_name() -> Arc>> { HW_ENCODER_NAME.clone() } pub fn encode(&mut self, bgra: &[u8]) -> ResultType> { match self.pixfmt { AVPixelFormat::AV_PIX_FMT_YUV420P => hw::hw_bgra_to_i420( self.encoder.ctx.width as _, self.encoder.ctx.height as _, &self.encoder.linesize, &self.encoder.offset, self.encoder.length, bgra, &mut self.yuv, ), AVPixelFormat::AV_PIX_FMT_NV12 => hw::hw_bgra_to_nv12( self.encoder.ctx.width as _, self.encoder.ctx.height as _, &self.encoder.linesize, &self.encoder.offset, self.encoder.length, bgra, &mut self.yuv, ), } match self.encoder.encode(&self.yuv) { Ok(v) => { let mut data = Vec::::new(); data.append(v); Ok(data) } Err(_) => Ok(Vec::::new()), } } } pub struct HwDecoder { decoder: Decoder, pub info: CodecInfo, } pub struct HwDecoders { pub h264: Option, pub h265: Option, } impl HwDecoder { /// See HwEncoder::best fn best(force_reset: bool, write: bool) -> (CodecInfos, bool) { let config = get_config(CFG_KEY_DECODER); if !force_reset && config.is_ok() { (config.unwrap(), false) } else { let decoders = CodecInfo::score(Decoder::avaliable_decoders()); if write { set_config(CFG_KEY_DECODER, &decoders) .map_err(|e| log::error!("{:?}", e)) .ok(); } (decoders, true) } } pub fn new_decoders() -> HwDecoders { let (best, _) = HwDecoder::best(false, true); let mut h264: Option = None; let mut h265: Option = None; let mut fail = false; if let Some(info) = best.h264 { h264 = HwDecoder::new(info).ok(); if h264.is_none() { fail = true; } } if let Some(info) = best.h265 { h265 = HwDecoder::new(info).ok(); if h265.is_none() { fail = true; } } if fail { HwDecoder::best(true, true); } HwDecoders { h264, h265 } } pub fn new(info: CodecInfo) -> ResultType { let ctx = DecodeContext { name: info.name.clone(), device_type: info.hwdevice.clone(), }; match Decoder::new(ctx) { Ok(decoder) => Ok(HwDecoder { decoder, info }), Err(_) => Err(anyhow!(format!("Failed to create decoder"))), } } pub fn decode(&mut self, data: &[u8]) -> ResultType> { match self.decoder.decode(data) { Ok(v) => Ok(v.iter().map(|f| HwDecoderImage { frame: f }).collect()), Err(_) => Ok(vec![]), } } } pub struct HwDecoderImage<'a> { frame: &'a DecodeFrame, } impl HwDecoderImage<'_> { pub fn bgra(&self, bgra: &mut Vec, i420: &mut Vec) -> ResultType<()> { let frame = self.frame; match frame.pixfmt { AVPixelFormat::AV_PIX_FMT_NV12 => hw::hw_nv12_to_bgra( frame.width as _, frame.height as _, &frame.data[0], &frame.data[1], frame.linesize[0] as _, frame.linesize[1] as _, bgra, i420, HW_STRIDE_ALIGN, ), AVPixelFormat::AV_PIX_FMT_YUV420P => { hw::hw_i420_to_bgra( frame.width as _, frame.height as _, &frame.data[0], &frame.data[1], &frame.data[2], frame.linesize[0] as _, frame.linesize[1] as _, frame.linesize[2] as _, bgra, ); return Ok(()); } } } } fn get_config(k: &str) -> ResultType { let v = HwCodecConfig::load() .options .get(k) .unwrap_or(&"".to_owned()) .to_owned(); match CodecInfos::deserialize(&v) { Ok(v) => Ok(v), Err(_) => Err(anyhow!("Failed to get config:{}", k)), } } fn set_config(k: &str, v: &CodecInfos) -> ResultType<()> { match v.serialize() { Ok(v) => { let mut config = HwCodecConfig::load(); config.options.insert(k.to_owned(), v); config.store(); Ok(()) } Err(_) => Err(anyhow!("Failed to set config:{}", k)), } } pub fn check_config() { let (encoders, update_encoders) = HwEncoder::best(false, false); let (decoders, update_decoders) = HwDecoder::best(false, false); if update_encoders || update_decoders { if let Ok(encoders) = encoders.serialize() { if let Ok(decoders) = decoders.serialize() { let mut config = HwCodecConfig::load(); config.options.insert(CFG_KEY_ENCODER.to_owned(), encoders); config.options.insert(CFG_KEY_DECODER.to_owned(), decoders); config.store(); return; } } log::error!("Failed to serialize codec info"); } }