matrix_thing/include/raii/video_man.hpp

/**
	This file is a part of rexy's matrix bot
	Copyright (C) 2019 rexy712

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU Affero General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Affero General Public License for more details.

	You should have received a copy of the GNU Affero General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef RAII_VIDEO_MAN_HPP
#define RAII_VIDEO_MAN_HPP

#include <memory> //unique_ptr
#include <utility> //exchange
#include "raii/string.hpp"
#include "raii/static_string.hpp"

//needs to be extern C because quality library writing
extern "C"{
#	include <libavcodec/avcodec.h> //AVCodecContext, AVCodec
#	include <libavformat/avformat.h> //AVFormatContext
#	include <libswscale/swscale.h> //sws_scale
#	include <libavutil/imgutils.h> //av_image_alloc
}

namespace raii{

	class video_man
	{
	private:
		static inline auto _context_deleter = [](AVCodecContext* ptr){
			avcodec_close(ptr);
			avcodec_free_context(&ptr);
		};
		static inline auto _frame_deleter = [](AVFrame* ptr){
			av_frame_unref(ptr);
			av_frame_free(&ptr);
		};
		static inline auto _packet_deleter = [](AVPacket* ptr){
			av_packet_unref(ptr);
			av_packet_free(&ptr);
		};
	public:
		using context_type = std::unique_ptr<AVCodecContext,decltype(_context_deleter)>;
		using frame_type = std::unique_ptr<AVFrame,decltype(_frame_deleter)>;
		using packet_type = std::unique_ptr<AVPacket,decltype(_packet_deleter)>;

	private:
		AVFormatContext* m_vid;
		AVCodec*         m_codec;
		AVCodecContext*  m_context;
		int              m_stream_index = -1;


	public:
		video_man(const raii::string_base& filename){
			av_register_all();
			_init(filename);
		}
		video_man(const video_man& v) = delete;
		video_man(video_man&& v):
			m_vid(std::exchange(v.m_vid, nullptr)){}
		~video_man(void){
			reset();
		}

		int height(void)const{
			return m_context->height;
		}
		int width(void)const{
			return m_context->width;
		}
		raii::string get_mimetype(void)const{
			const char* first = strstr(m_vid->iformat->name, ",");
			return raii::string(raii::static_string(m_vid->iformat->name, first - m_vid->iformat->name));
		}
		packet_type create_jpg_thumbnail(void){
			context_type jpg_context = _get_jpeg_context(m_context);
			frame_type frame(_init_frame(m_context->width, m_context->height, m_context->pix_fmt), _frame_deleter);
			frame_type jpg_frame(_init_frame(m_context->width, m_context->height, jpg_context->pix_fmt), _frame_deleter);
			packet_type packet(av_packet_alloc(), _packet_deleter);
			SwsContext* sws_ctx = sws_getContext(m_context->width, m_context->height, m_context->pix_fmt,
			                                     m_context->width, m_context->height, AV_PIX_FMT_YUV420P,
			                                     SWS_BILINEAR, NULL, NULL, NULL);
			_change_color_range(sws_ctx);

			av_image_alloc(jpg_frame->data, jpg_frame->linesize, m_context->width, m_context->height, jpg_context->pix_fmt, 32);
			while(av_read_frame(m_vid, packet.get()) >= 0){
				if(packet->stream_index == m_stream_index){
					avcodec_send_packet(m_context, packet.get());
					if(avcodec_receive_frame(m_context, frame.get()) >= 0)
						break;
				}
				av_packet_unref(packet.get());
			}
			sws_scale(sws_ctx, frame->data, frame->linesize, 0, m_context->height, jpg_frame->data, jpg_frame->linesize);

			sws_freeContext(sws_ctx);

			packet_type jpg_packet(av_packet_alloc(), _packet_deleter);
			if(avcodec_send_frame(jpg_context.get(), jpg_frame.get()) < 0){
				av_freep(&jpg_frame->data[0]);
				return packet_type(nullptr, _packet_deleter);
			}
			if(avcodec_receive_packet(jpg_context.get(), jpg_packet.get()) < 0){
				return packet_type(nullptr, _packet_deleter);
			}

			av_freep(&jpg_frame->data[0]);
			return jpg_packet;
		}

		AVFormatContext* release(void){
			avcodec_close(m_context);
			avcodec_free_context(&m_context);
			m_context = nullptr;
			m_codec = nullptr;
			m_stream_index = -1;
			return std::exchange(m_vid, nullptr);
		}
		void reset(const raii::string_base& filename){
			reset();
			_init(filename);
		}
		void reset(void){
			avcodec_close(m_context);
			avcodec_free_context(&m_context);
			avformat_close_input(&m_vid);
			m_context = nullptr;
			m_vid = nullptr;
			m_codec = nullptr;
			m_stream_index = -1;
		}
	private:
		void _init(const raii::string_base& filename){
			m_vid = _open_video_file(filename);
			if(!m_vid)
				return;

			AVCodecParameters* codec_parms = m_vid->streams[m_stream_index]->codecpar;
			m_codec = avcodec_find_decoder(codec_parms->codec_id);
			m_context = avcodec_alloc_context3(m_codec);
			avcodec_parameters_to_context(m_context, codec_parms);
			avcodec_open2(m_context, m_codec, NULL);
		}
		AVFormatContext* _open_video_file(const raii::string_base& filename){
			AVFormatContext* context = NULL;
			avformat_open_input(&context, filename.get(), NULL, NULL);
			avformat_find_stream_info(context, NULL);
			for(size_t i = 0;i < context->nb_streams;++i){
				if(context->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO){
					m_stream_index = i;
					break;
				}
			}
			if(m_stream_index != -1)
				return context;
			avformat_close_input(&context);
			return nullptr;
		}
		static AVFrame* _init_frame(int w, int h, int format){
			AVFrame* f = av_frame_alloc();
			f->data[0] = NULL;
			f->width = w;
			f->height = h;
			f->format = format;
			return f;
		}
		static void _change_color_range(SwsContext* sws_ctx){
			//we want to convert from YUVJ420P to YUV420P for sws_scale to be happy. but we actually want to output as YUVJ420P.
			//so we change the color range to use the full range like YUVJ420P. dummy is used to ignore unneeded returns.
			int dummy[4];
			int src_range, dest_range;
			int br, con, sat;
			sws_getColorspaceDetails(sws_ctx, (int**)&dummy, &src_range, (int**)&dummy, &dest_range, &br, &con, &sat);
			const int* coefs = sws_getCoefficients(SWS_CS_DEFAULT);
			src_range = 2;
			sws_setColorspaceDetails(sws_ctx, coefs, src_range, coefs, dest_range, br, con, sat);
		}
		static context_type _get_jpeg_context(const AVCodecContext* decoder){
			AVCodec* jpg_codec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
			context_type jpg_context(avcodec_alloc_context3(jpg_codec), _context_deleter);
			jpg_context->pix_fmt = AV_PIX_FMT_YUVJ420P;
			jpg_context->width = decoder->width;
			jpg_context->height = decoder->height;
			jpg_context->codec_id = AV_CODEC_ID_MJPEG;
			jpg_context->codec_type = AVMEDIA_TYPE_VIDEO;
			if(decoder->time_base.num == 0 || decoder->time_base.den == 0)
				jpg_context->time_base = AVRational{1, 30};
			else
				jpg_context->time_base = decoder->time_base;
			avcodec_open2(jpg_context.get(), jpg_codec, NULL);
			return jpg_context;
		}
	};


}

#endif