Making Screenshots at Runtime

#pragma once

#include <UnigineComponentSystem.h>
#include <UnigineViewport.h>

class ScreenshotMaker : public Unigine::ComponentBase
{
	enum Format
	{
		tga = 0,
		png,
		jpg
	};

	enum State
	{
		STATE_START = 0,
		STATE_WARMUP,
		STATE_SAVE,
		STATE_DONE,
	};

public:
	COMPONENT_DEFINE(ScreenshotMaker, Unigine::ComponentBase);

	PROP_PARAM(String, name_prefix, "screenshot");
	PROP_PARAM(IVec2, size, Unigine::Math::ivec2(640, 480));
	PROP_PARAM(Switch, format, 0, "tga,png,jpg");
	PROP_PARAM(Toggle, alfa_channel, false);

	COMPONENT_INIT(init);
	COMPONENT_UPDATE(update);

private:
	Unigine::ImagePtr image;
	Unigine::ViewportPtr viewport;
	int count = 0;

	int state = STATE_DONE;
	const int warmup_frames = 30;
	int warm_up_count = 0;

	void init();
	void update();
	
};

#include "ScreenshotMaker.h"
#include <UnigineConsole.h>
#include <UnigineInput.h>
#include <UnigineApp.h>
#include <UnigineGame.h>

REGISTER_COMPONENT(ScreenshotMaker);

using namespace Unigine;

void ScreenshotMaker::init()
{
	image = Image::create();
	viewport = Viewport::create();
	viewport->setSkipFlags(Viewport::SKIP_VISUALIZER);

	Console::setOnscreen(1);
	Console::onscreenMessageLine(Math::vec4_green, "Screenshot component is initialized.");
}

void ScreenshotMaker::update()
{
	auto player = Game::getPlayer();
	if (!player)
	{
		state = STATE_DONE;
		Console::onscreenMessageLine(Math::vec4_red, "No active camera.");
		return;
	}

	if (state == STATE_DONE && App::clearKeyState('t'))
	{
		warm_up_count = 0;
		state = STATE_WARMUP;
		Console::onscreenMessageLine(Math::vec4_green, "Trying to take a screenshot...");
	}

	if (state == STATE_WARMUP)
	{
		viewport->setMode(Render::getViewportMode());

		if (warm_up_count == 0)
		{
			// First frame we render with velocity buffer turned off to avoid temporal effects artifacts
			viewport->appendSkipFlags(Viewport::SKIP_VELOCITY_BUFFER);
			viewport->renderImage2D(player->getCamera(), image, size.get().x, size.get().y);
			viewport->removeSkipFlags(Viewport::SKIP_VELOCITY_BUFFER);
		}
		else
		{
			// We temporarily set exposure adaptation time to 0, otherwise the image may be too dark
			float exposure_adaptation = Render::getExposureAdaptation();
			Render::setExposureAdaptation(0.0f);
			viewport->renderImage2D(player->getCamera(), image, size.get().x, size.get().y);
			Render::setExposureAdaptation(exposure_adaptation);
		}

		warm_up_count++;
		if (warm_up_count == warmup_frames)
			state = STATE_SAVE;
	}

	if (state == STATE_SAVE)
	{
		if (!alfa_channel.get() || format.get() == Format::jpg)
		{
			if (image->getFormat() == Image::FORMAT_RGBA8)
				image->convertToFormat(Image::FORMAT_RGB8);
			else if (image->getFormat() == Image::FORMAT_RGBA16F)
				image->convertToFormat(Image::FORMAT_RGB16F);
		}

		auto str_formats = String::split("tga,png,jpg", ",");

		String fullName = String::format("%s_%d.%s", name_prefix.get(), count, str_formats[format.get()]);
		image->save(fullName.get());
		Console::onscreenMessageLine(Math::vec4_green, "%s saved.", fullName.get());

		count++;
		state = STATE_DONE;
	}
}

int AppSystemLogic::init()
{
	Unigine::ComponentSystem::get()->initialize();
	// Write here code to be called on engine initialization.
	return 1;
}