mozaic/src/main.cpp

#include <image_io.hpp>
#include <wavelet.hpp>
#include <image_database.hpp>
#include <fingerprints.hpp>
#include <mozaic.hpp>

#include <boost/filesystem.hpp>
#include <boost/archive/binary_oarchive.hpp>
#include <boost/archive/binary_iarchive.hpp>

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
}

#include <iostream>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <map>
#include <set>

static const int tile_width = 128;
static const int tile_height = 128;

using namespace std;
namespace fs = boost::filesystem;
namespace ar = boost::archive;

using Database = image_database<downscale>;
using Mozaiq = Mozaic<string>;

static Database read_database(string const & database_directory){
	Database db;
	auto const database_file = database_directory + ".db";

	if (!boost::filesystem::exists(database_file)){
		fs::path const directory(database_directory);
		fs::directory_iterator eod;
		for(fs::directory_iterator it(directory); it != eod; ++it){
			auto const path = it->path();
			auto const ext = path.extension();
			if(ext != ".png" && ext != ".jpg") continue;

			cout << colors::green("adding: ") << path.string() << endl;
			db.add(path.string());
		}

		ofstream file(database_file);
		ar::binary_oarchive archive(file);
		archive << db;
	} else {
		ifstream file(database_file);
		ar::binary_iarchive archive(file);
		archive >> db;
	}

	return db;
}

static auto calculate_all_distances(Database const & db, string const & filename, int h_tiles, int v_tiles){
	using return_type = decltype(db.distances_for_image(std::declval<av::frame const &>()));
	Mozaic<return_type> mozaic(h_tiles, v_tiles);

	apply_to_tiles(filename, h_tiles, v_tiles, [&](int c, int r, av::frame const & frame){
		mozaic[r][c] = db.distances_for_image(frame);
	});

	return mozaic;
}

static void save_mozaic(Mozaiq const & mozaic, string filename){
	auto const pix_fmt = AV_PIX_FMT_YUVJ444P;
	auto const codec_id= AV_CODEC_ID_MJPEG;

	// Open all files we need
	map<string, av::frame> frames;
	for(auto&& y : mozaic.tiles){
		for(auto&& x : y){
			// only open a file once
			if(frames.count(x) > 0) continue;
			frames.emplace(x, crop_to_square(open_image(x)));
		}
	}

	auto const total_width = mozaic.h_tiles * tile_width;
	auto const total_height = mozaic.v_tiles * tile_height;

	// Create output frame
	std::vector<uint8_t, av::allocator<uint8_t>> data(make_u(avpicture_get_size(pix_fmt, total_width, total_height)), 0);
	av::frame frame = av::frame_alloc();
	avpicture_fill(reinterpret_cast<AVPicture*>(frame.get()), data.data(), pix_fmt, total_width, total_height);
	frame->width = total_width;
	frame->height = total_height;
	frame->format = pix_fmt;

	// For each tile: get the part, copy input to it
	av::frame frame_part = av::frame_clone(frame);
	for(int r = 0; r < mozaic.v_tiles; ++r) {
		for(int c = 0; c < mozaic.h_tiles; ++c){
			av_picture_crop(reinterpret_cast<AVPicture*>(frame_part.get()), reinterpret_cast<AVPicture*>(frame.get()), av::get_format(frame), r * tile_height, c * tile_width);
			frame_part->width = tile_width;
			frame_part->height = tile_height;

			auto&& input = frames.at(mozaic[r][c]);
			auto sws_context = sws_getContext(input->width, input->height, av::get_format(input), frame_part->width, frame_part->height, av::get_format(frame_part), 0, nullptr, nullptr, nullptr);
			if(!sws_context) throw std::runtime_error("boem sws context");
			sws_scale (sws_context, {input->data}, {input->linesize}, 0, input->height, {frame_part->data}, {frame_part->linesize});
			sws_freeContext(sws_context);
		}
	}

	// Done with the input
	frames.clear();

	// Encode
	auto codec = avcodec_find_encoder(codec_id);
	if(!codec) throw av::error("Could not find codec");

	auto codec_ctx = std::unique_ptr<AVCodecContext, av::deleter<AVCodecContext>>(avcodec_alloc_context3(codec), [](auto x){ avcodec_free_context(&x); });
	if(!codec_ctx) throw av::error("Could not allocate codec context");

	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);
	auto opened_codec = av::codec_open(codec_ctx.get(), codec, nullptr);

	auto const buffer_size = avpicture_get_size(pix_fmt, codec_ctx->width, codec_ctx->height);
	std::vector<uint8_t> buffer(make_u(buffer_size), 0);
	auto const output_size = avcodec_encode_video(codec_ctx.get(), buffer.data(), buffer_size, frame.get());
	assert(output_size <= buffer_size);

	auto file = fopen(filename.c_str(), "wb");
	fwrite(buffer.data(), 1, make_u(output_size), file);
	fclose(file);
}

int main(){
	av_register_all();
	av_log_set_level(AV_LOG_QUIET);

	// TODO: use boost::program_options
	string const database_directory = "database";
	string const filename = "image.jpg";
	string const output = "output.jpg";
	int const h_tiles = 4 * 7;
	int const v_tiles = 3 * 7;

	// NOTE: we could almost do the rest on one line ;D
	// TODO: make every step parallel (this should be easy, except for the IO steps)
	auto const db = read_database(database_directory);
	cout << colors::green("Read database: ") << db.size() << endl;

	auto const distances = calculate_all_distances(db, filename, h_tiles, v_tiles);
	cout << colors::green("Calculated distances for ") << filename << endl;

	auto const best_indices = optimize(distances, 1);
	cout << colors::green("Picked best distances") << endl;

	auto const mozaic = mozaic_fmap(best_indices, [&](auto i){ return db.filename(i); });
	cout << colors::green("Picked filenames") << endl;

	save_mozaic(mozaic, output);
	cout << colors::green("Saved Mozaic to ") << output << endl;
}