thumbnail.c

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>

#include "util.h"

AVFrame *scale_frame(AVFrame *frame, enum AVPixelFormat pix_fmt,
		int width, int height)
{
	AVFrame *scaled;
	struct SwsContext *swsctx;
	scaled = av_frame_alloc();
	scaled->format = pix_fmt;
	scaled->width = width;
	scaled->height = height;
	av_frame_get_buffer(scaled, 32);
	swsctx = sws_getContext(frame->width, frame->height, frame->format,
			scaled->width, scaled->height, scaled->format, 0, 0, 0, 0);
	sws_scale(swsctx, (const uint8_t **) frame->data, frame->linesize, 0,
			frame->height, scaled->data, scaled->linesize);
	return scaled;
}

int main(int argc, char *argv[])
{
	int err, vidx;
	char *src, *tmp, *dest;
	AVFormatContext *demux, *mux;
	AVCodecContext *dec, *enc;
	AVPacket pkt;
	AVFrame *frame;
	if (argc != 5)
		return 1;
	av_register_all();
	avformat_network_init();
	init_env();
	init_avlogbuf();

	/* Prepare arguments */
	jobid = atoi(argv[1]);
	ping_job(jobid, "running", 0);
	if (jlookup(argv[4], "srcurl"))
		src = jstr(jlookup(argv[4], "srcurl"), 0);
	else
		src = buildpath(getenv(WORKER_RELEASED), jstr(jlookup(argv[4], "src"), 0));
	tmp = mprintf("%s/.tmp-%i", getenv(WORKER_TMP), jobid);
	dest = mprintf("%s/thumbnail/%s", getenv(WORKER_RELEASED),
			jstr(jlookup(argv[4], "filename"), "0"));

	/* Open src */
	demux = 0;
	if (err = avformat_open_input(&demux, src, 0, 0))
		job_failed("Opening input file failed: %s", av_err2str(err));
	avformat_find_stream_info(demux, 0);
	if ((vidx = av_find_best_stream(demux, AVMEDIA_TYPE_VIDEO, -1, -1, 0, 0)) < 0)
		job_failed("Could not find video stream in input file");
	dec = avcodec_alloc_context3(avcodec_find_decoder(demux->streams[vidx]->codecpar->codec_id));
	avcodec_parameters_to_context(dec, demux->streams[vidx]->codecpar);
	avcodec_open2(dec, dec->codec, 0);

	/* Open temporary dest */
	avformat_alloc_output_context2(&mux, 0, "image2", tmp);
	avformat_new_stream(mux, 0);
	enc = avcodec_alloc_context3(avcodec_find_encoder(AV_CODEC_ID_MJPEG));
	enc->width = 640;
	enc->height = (enc->width*dec->height)/dec->width;
	if (enc->codec->pix_fmts)
		enc->pix_fmt = enc->codec->pix_fmts[0];
	else
		enc->pix_fmt = dec->pix_fmt;
	enc->sample_aspect_ratio = dec->sample_aspect_ratio;
	enc->time_base = AV_TIME_BASE_Q;
	avio_open(&mux->pb, tmp, AVIO_FLAG_WRITE);
	avcodec_open2(enc, enc->codec, 0);
	avcodec_parameters_from_context(mux->streams[0]->codecpar, enc);
	if ((err = avformat_write_header(mux, 0)) < 0)
		job_failed("Writing temporary file failed: %s", av_err2str(err));

	/* Create thumbnail */
	if (!jlookup(argv[4], "srcurl"))
		av_seek_frame(demux, -1, (2L*demux->duration)/5L, 0);
	memset(&pkt, 0, sizeof(pkt));
	av_init_packet(&pkt);
	frame = av_frame_alloc();
	do
	{
		if (err = av_read_frame(demux, &pkt))
			job_failed("Reading frame failed: %s", av_err2str(err));
		if (pkt.stream_index == vidx)
			avcodec_send_packet(dec, &pkt);
	}
	while (avcodec_receive_frame(dec, frame));
	avcodec_send_frame(enc, scale_frame(frame, enc->pix_fmt, enc->width,
				enc->height));
	avcodec_send_frame(enc, 0);
	while (!avcodec_receive_packet(enc, &pkt))
	{
		pkt.stream_index = 0;
		pkt.pts = pkt.dts = 0;
		av_interleaved_write_frame(mux, &pkt);
	}
	av_interleaved_write_frame(mux, 0);
	if (err = av_write_trailer(mux))
		job_failed("Error writing trailer to temporary file: %s", av_err2str(err));
	avio_closep(&mux->pb);
	if (!filesize(tmp))
		job_failed("Sanity check failed: Output file is empty");
	if (rename(tmp, dest))
		job_failed("Overwriting output file failed: %s", strerror(errno));
	unlink(tmp);
	ping_job(jobid, "finished", "{\"log\": \"%s\"}", jescape(get_avlogbuf()));
	return 0;
}