CaptureSystem.cpp

#include "CaptureSystem.h"
 
#include "Context.h"
#include "Services/Time.h"
#include "Components/Core/CameraComponent.h"
#include "Systems/Core/CameraSystem.h"
#include "Services/Time.h"
#include "Resources/TextureManager.h"
#include "Renderer/RenderResource.h"
 
#include "Foundation/Logging/Logger.h"
#include "Foundation/Platform/IO.h"
 
#include <chrono>
 
#ifndef EMSCRIPTEN
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "../../../Libraries/stb/stb_image_write.h"
#endif
 
using namespace Cogs::Core;
 
namespace
{
  using namespace Cogs::Core;
 
  Cogs::Logging::Log logger = Cogs::Logging::getLogger("CaptureSystem");
 
#ifndef EMSCRIPTEN
  void writeFunc(void *context, void *data, int size)
  {
    // Note: stbi_write_png_to_func creates an intermediate buffer
    // that is passed to this callback and is finished with the
    // input data when it is called, so we can overwrite the input.
    auto *v = reinterpret_cast<Cogs::Memory::MemoryBuffer* > (context);
    v->resize(size, false);
    std::memcpy(v->data(), data, size);
  }
#endif
 
  struct PNGEncodingTask
  {
    std::atomic<CaptureItem*>& taskItem;
 
    void operator()()
    {
      auto * item = taskItem.load();
 
      using namespace std::chrono_literals;
      //LOG_DEBUG(logger, "PNG encoding of frame %d", item->frame);
 
#ifdef EMSCRIPTEN
      item->state = CaptureItem::State::Dropped;
#else
 
      if (stbi_write_png_to_func(writeFunc, &item->data, item->width, item->height, 4, item->data.data(), 4 * item->width) == 0)
      {
        LOG_WARNING(logger, "PNG encoding of frame %d failed", item->frame);
        item->state = CaptureItem::State::Dropped;
      }
      else {
        item->state = CaptureItem::State::EncodedAsPng;
      }
#endif
      //std::this_thread::sleep_for(10ms);
 
      taskItem.store(nullptr);
    }
  };
 
  struct StoreImageTask
  {
    std::atomic<CaptureItem*>& taskItem;
 
    void operator()()
    {
      auto * item = taskItem.load();
#ifdef EMSCRIPTEN
      LOG_WARNING(logger, "Storing of images not implemented.");
#else
 
      char sep[2];
      if (item->recordPath.empty()) {
        sep[0] = '\0';
      }
      else {
        sep[0] = Cogs::IO::pathSeparator();
        sep[1] = '\0';
      }
 
      auto seconds = int(std::floor(item->time));
      auto millis = int(std::floor(1000.f*(item->time - seconds)));
 
      const char* suffix = nullptr;
      switch (item->format) {
      case CaptureFormat::PNG:
        suffix = "png";
        break;
      case CaptureFormat::RawARGB:
        suffix = "rgba";
        break;
      case CaptureFormat::RawFloat:
        suffix = "float";
        break;
      default:
        assert(false);
      }
 
      char buffer[1024];
      snprintf(buffer, sizeof(buffer), "%s%s%s-f%05d-t%05d_%03d-%dx%d.%s", item->recordPath.c_str(), sep, item->name.c_str(), item->frame, seconds, millis, item->width, item->height, suffix);
 
 
      LOG_DEBUG(logger, "Writing '%s'", buffer);
      Cogs::IO::writeBinaryFile(buffer, item->data.data(), item->data.size());
#endif
      item->state = CaptureItem::State::Ready;
      taskItem.store(nullptr);
    }
 
 
  };
 
  void readbackCallback(Context * context, const CameraData* camData, uint32_t frame, const char * key, const void * data, int size_)
  {
    const auto size = static_cast<unsigned>(size_);
    if (std::string(key) != "Capture") return;
 
    auto * capComp = camData->camera.resolve()->getComponent<CaptureComponent>();
    if (!capComp) {
      LOG_FATAL(logger, "No capture component attached to camera.");
      return;
    }
    auto & capData = context->captureSystem->getData(capComp);
 
    for (auto & item : capData.items) {
      if (item->frame == frame) {
 
        item->width = int(std::max(0.f, camData->viewportSize[0]));
        item->height = int(std::max(0.f, camData->viewportSize[1]));
 
        if (item->format == CaptureFormat::PNG) {
 
          auto expectedSize = 4 * item->width * item->height;
          if (size != expectedSize) {
            LOG_WARNING(logger, "Expected %d bytes of image data, got %d bytes, dropping frame %d", expectedSize, size, frame);
            return;
          }
 
          if (capData.encodePNGItem.load()) {
            if (capComp->dropFrames) {
              //LOG_WARNING(logger, "PNG encoding of previous frame not finished, dropping frame %d.", frame);
              item->state = CaptureItem::State::Dropped;
              return;
            }
            else {
              context->taskManager->wait(capData.encodePngTask);
              assert(capData.encodePNGItem.load() == nullptr);
              capData.encodePngTask = NoTask;
            }
          }
 
          item->state = CaptureItem::State::EncodeToPng;
          item->data.resize(size, false);
          std::memcpy(item->data.data(), data, size);
 
          capData.encodePNGItem.store(item.get());
          capData.encodePngTask = context->taskManager->enqueue(context->taskManager->GlobalQueue,
                                                          PNGEncodingTask{ capData.encodePNGItem });
          //LOG_DEBUG(logger, "Scheduled frame %d for PNG encoding", frame);
        }
        else {
          item->state = CaptureItem::State::ReceivedReadback;
          item->data.resize(size, false);
          std::memcpy(item->data.data(), data, size);
          //LOG_DEBUG(logger, "Received data for frame %d", frame);
        }
        if (0 < capComp->framesToCapture) {
          capComp->framesToCapture--;
        }
        return;
      }
    }
    //LOG_WARNING(logger, "Got capture readback data but couldn't find appropriate metadata item for frame %d.", frame)
  }
 
 
  void removeDroppedItems(std::vector<std::unique_ptr<CaptureItem>>& itemStore, std::deque<std::unique_ptr<CaptureItem>>& items)
  {
    for (auto it = items.begin(); it != items.end(); ) {
      if ((*it)->state == CaptureItem::State::Dropped) {
        //LOG_DEBUG(logger, "Discarded dropped item");
        itemStore.push_back(std::move(*it));
        it = items.erase(it);
      }
      else {
        ++it;
      }
    }
  }
 
  void removeScheduledForDiscard(std::vector<std::unique_ptr<CaptureItem>>& itemStore, std::deque<std::unique_ptr<CaptureItem>>& items)
  {
    // Check if there is anything besides just issued frames,
    // in that case we will delete all issued frames up until the first that we have data on.
    bool anyNotJustIssued = false;
    for (auto & item : items) {
      if (item->state != CaptureItem::State::Issued) {
        anyNotJustIssued = true;
        break;
      }
    }
 
    for (auto it = items.begin(); it != items.end(); ) {
      auto & item = *it;
      bool discard = false;
 
      if (anyNotJustIssued && item->state == CaptureItem::State::Issued) {
        discard = true;
      }
      else {
        anyNotJustIssued = false;
        if (item->state == CaptureItem::State::DiscardNextFrame) {
          discard = true;
        }
      }
 
      if (discard) {
        itemStore.push_back(std::move(*it));
        it = items.erase(it);
      }
      else {
        ++it;
      }
    }
 
 
 
 
    for (auto & item : items) {
      if (item->state == CaptureItem::State::Ready) {
        item->state = CaptureItem::State::DiscardNextFrame;
      }
    }
 
  }
 
 
}
 
void Cogs::Core::CaptureSystem::initialize(Context * context)
{
  ComponentSystemWithDataPool::initialize(context);
  context->callbacks.readbackCallbackInternal.push_back(readbackCallback);
}
 
void Cogs::Core::CaptureSystem::update(Context * context)
{
  for (auto & capComp : pool) {
    auto & capData = getData(&capComp);
 
    auto * camComp = capComp.getComponent<CameraComponent>();
    if (!camComp) continue;
 
    bool doCapture = capComp.framesToCapture != 0;
    
    char args[256];
    snprintf(args, sizeof(args), "Capture=%s&Format=%s", doCapture ? "True" : "False", capComp.format == CaptureFormat::RawFloat ? "Float" : "4xU8");
 
    switch (capComp.mode) {
    case CaptureMode::Off:
      camComp->flags &= ~CameraFlags::EnableRender;
      camComp->renderPipeline = "";
      break;
    case CaptureMode::Scene:
      camComp->flags |= CameraFlags::EnableRender;
      switch (capComp.sceneMode) {
      case CaptureSceneMode::Forward:
        camComp->renderPipeline = "Pipelines/CaptureForward.pipeline?" + std::string(args);
        break;
      case CaptureSceneMode::ForwardHdr:
        camComp->renderPipeline = "Pipelines/CaptureForwardHdr.pipeline?" + std::string(args);
        break;
      case CaptureSceneMode::DeferredCausticsHdr:
        camComp->renderPipeline = "Pipelines/CaptureDeferredCausticsHdr.pipeline?" + std::string(args);
        break;
      default:
        assert(false);
      }
      break;
    case CaptureMode::ObjectId:
      camComp->flags |= CameraFlags::EnableRender;
      camComp->renderPipeline = "Pipelines/Capture.pipeline?Mode=ObjectId&" + std::string(args);
      break;
    case CaptureMode::ObjectIdColorize:
      camComp->flags |= CameraFlags::EnableRender;
      camComp->renderPipeline = "Pipelines/Capture.pipeline?Mode=ObjectIdColorize&" + std::string(args);
      break;
    case CaptureMode::Normals:
      camComp->flags |= CameraFlags::EnableRender;
      camComp->renderPipeline = "Pipelines/Capture.pipeline?Mode=Normals&" + std::string(args);
      break;
    case CaptureMode::Depth:
      camComp->flags |= CameraFlags::EnableRender;
      camComp->renderPipeline = "Pipelines/Capture.pipeline?Mode=Depth&" + std::string(args);
      break;
    default:
      assert(false);
    }
    camComp->setChanged();
 
    for (auto & item : capData.items) {
      if (item->state == CaptureItem::State::ReceivedReadback || item->state == CaptureItem::State::EncodedAsPng)
      {
        if (item->recordToDisc) {
          if (capData.storeImageItem.load() && capComp.dropFrames) {
            LOG_WARNING(logger, "Storing of previous frame not finished, skip storing of frame %d.", item->frame);
            item->state = CaptureItem::State::Ready;
          }
          else {
            if (capData.storeImageTask.isValid()) {
              context->taskManager->wait(capData.storeImageTask);
              capData.storeImageTask = NoTask;
            }
            assert(capData.storeImageItem.load() == nullptr);
            item->state = CaptureItem::State::WriteToDisc;
            capData.storeImageItem.store(item.get());
            assert(capData.storeImageItem.load());
            capData.storeImageTask = context->taskManager->enqueue(context->taskManager->GlobalQueue, StoreImageTask{ capData.storeImageItem });
          }
        }
        else {
          item->state = CaptureItem::State::Ready;
        }
      }
    }
 
    removeDroppedItems(capData.itemStore, capData.items);
    removeScheduledForDiscard(capData.itemStore, capData.items);
 
    if (doCapture) {
      if (10 <= capData.items.size()) {
        if (capData.items.front()->state == CaptureItem::State::Issued) {
          capData.items.pop_front();
        }
        else {
          LOG_WARNING(logger, "10 or more capture frames in flight for CaptureComponent, skipping.");
        }
      }
      else {
        auto * entity = capComp.getContainer();
        std::unique_ptr<CaptureItem> item;
        if (capData.itemStore.empty()) {
          item = std::make_unique<CaptureItem>();
        }
        else {
          item = std::move(capData.itemStore.back());
          capData.itemStore.pop_back();
        }
        item->time = context->time->getAnimationTime();
        item->frame = context->time->getFrame();
        item->state = CaptureItem::State::Issued;
        item->mode = capComp.mode;
        item->sceneMode = capComp.sceneMode;
        item->format = capComp.format;
        item->recordToDisc = capComp.recordToDisc;
        item->recordPath = capComp.recordPath;
        item->name = entity->getName();
        if (item->name.empty()) {
          item->name = "entity_" + std::to_string(entity->getId());
        }
        capData.items.push_back(std::move(item));
      }
    }
  }
}
 
 
void Cogs::Core::CaptureSystem::destroyComponent(ComponentHandle component)
{
  auto * comp = component.resolveComponent<CaptureComponent>();
  auto & data = getData(comp);
  if (data.encodePngTask.isValid()) {
    context->taskManager->wait(data.encodePngTask);
  }
  if (data.storeImageTask.isValid()) {
    context->taskManager->wait(data.storeImageTask);
  }
  ComponentSystemWithDataPool::destroyComponent(component);
}