Making mosaic images with ffmpeg
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

171 lines
5.0 KiB

#include "image_io.hpp"
#include "utilities.hpp"
#include <av/sws.hpp>
extern "C" {
#include <libavutil/frame.h>
#include <libavutil/mem.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
}
#include <algorithm>
#include <cassert>
#include <fstream>
#include <iostream>
#include <stdexcept>
raw_rgb_image::raw_rgb_image()
: data()
, frame(nullptr, [](auto x){ av_frame_free(&x); })
{}
raw_rgb_image::raw_rgb_image(int W, int H)
: data(make_u(avpicture_get_size(AV_PIX_FMT_RGB24, W, H)))
, frame(av::frame_alloc())
{
avpicture_fill(reinterpret_cast<AVPicture*>(frame.get()), data.data(), AV_PIX_FMT_RGB24, W, H);
frame->width = W;
frame->height = H;
frame->format = AV_PIX_FMT_RGB24;
}
int raw_rgb_image::width() const { return frame->width; }
int raw_rgb_image::height() const { return frame->height; }
AVPixelFormat raw_rgb_image::format() const { return av::get_format(frame); }
av::frame open_image(std::string const & filename){
// Open the file
auto format_context = av::format_open_input(filename, nullptr, nullptr);
// Get the codec and let us own the buffers
auto codec_context = av::context_from_stream(format_context, 0);
const auto codec = avcodec_find_decoder(codec_context->codec_id);
codec_context->refcounted_frames = 1;
// Open the codec
const auto opened_codec = av::codec_open(codec_context, codec, nullptr);
// Allocate frame
av::frame frame = av::frame_alloc();
// things to read and decode it
av::packet_buffer empty_packet;
int finished = 0;
while(auto packet = av::read_frame(format_context, empty_packet)) {
if(packet->stream_index != 0) continue;
int ret = avcodec_decode_video2(codec_context.get(), frame.get(), &finished, packet.get());
if (ret <= 0) {
printf("Error [%d] while decoding frame: %s\n", ret, strerror(AVERROR(ret)));
throw std::runtime_error("boem packet");
}
// we only need the first frame
if(finished) break;
}
// some decoders need extra passes
while(!finished) {
avcodec_decode_video2(codec_context.get(), frame.get(), &finished, &empty_packet);
}
return frame;
}
void crop_to_square(av::frame& frame){
int diff = frame->width - frame->height;
if(diff > 0) {
av::crop(frame, diff/2, 0, frame->height, frame->height);
} else if(diff < 0) {
av::crop(frame, 0, -diff/2, frame->width, frame->width);
}
}
av::frame crop_to_square(const av::frame& frame){
return crop_to_square(av::frame_clone(frame));
}
av::frame crop_to_square(av::frame && frame){
crop_to_square(frame);
return std::move(frame);
}
raw_rgb_image to_raw_rgb_image(av::frame const & frame, int new_width, int new_height){
raw_rgb_image image(new_width, new_height);
auto context = sws::create_context(frame, image.frame);
sws::scale(context, frame, image.frame);
return image;
}
void apply_to_tiles(std::string const & filename, int h_tiles, int v_tiles, std::function<void(int, int, av::frame const &)> fun) {
const auto org_frame = open_image(filename);
// create raw buffer for the callback
// TODO: do not scale the cropped region
raw_rgb_image image(512, 512);
// create the tiles
const int width = org_frame->width / h_tiles;
const int height = org_frame->height / v_tiles;
for(int r = 0; r < v_tiles; ++r){
for(int c = 0; c < h_tiles; ++c){
const int x_crop = c * width;
const int y_crop = r * height;
const auto cropped_frame = av::crop(org_frame, x_crop, y_crop, width, height);
auto context = sws::create_context(cropped_frame, image.frame);
sws::scale(context, cropped_frame, image.frame);
fun(c, r, image.frame);
}
}
}
void save_as_jpg(av::frame const & frame, std::string const & filename){
const auto pix_fmt = AV_PIX_FMT_YUVJ444P;
std::vector<uint8_t, av::allocator<uint8_t>> data(make_u(avpicture_get_size(pix_fmt, frame->width,
frame->height)), 0);
av::frame converted_frame = av::frame_alloc();
avpicture_fill(reinterpret_cast<AVPicture*>(converted_frame.get()), data.data(), pix_fmt, frame->width, frame->height);
converted_frame->width = frame->width;
converted_frame->height = frame->height;
converted_frame->format = pix_fmt;
{
auto sws_context = sws::create_context(frame, converted_frame);
sws::scale(sws_context, frame, converted_frame);
}
encode_as_jpg(converted_frame, filename);
}
void encode_as_jpg(const av::frame& frame, const std::string& filename){
auto const codec_id= AV_CODEC_ID_MJPEG;
const auto pix_fmt = av::get_format(frame);
const auto codec = av::find_encoder(codec_id);
auto codec_ctx = av::context_alloc(codec);
codec_ctx->pix_fmt = pix_fmt;
codec_ctx->width = frame->width;
codec_ctx->height = frame->height;
codec_ctx->time_base = av_make_q(1, 1);
const auto opened_codec = av::codec_open(codec_ctx, codec, nullptr);
const auto buffer_size = avpicture_get_size(pix_fmt, codec_ctx->width, codec_ctx->height);
std::vector<uint8_t> buffer(make_u(buffer_size), 0);
const auto output_size = avcodec_encode_video(codec_ctx.get(), buffer.data(), buffer_size, frame.get());
assert(output_size <= buffer_size);
std::ofstream file(filename);
file.write(reinterpret_cast<char*>(buffer.data()), output_size);
}