From ad5bee73ef1386f2448a6894a869e760355da11d Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Wed, 1 May 2024 08:58:12 -0700 Subject: [PATCH] Cache decoded PNG images for hfields and textures in the MuJoCo compiler. PiperOrigin-RevId: 629738092 Change-Id: Ie1b69a5dcee0695aadb2f02ef3b4d7a73cdf07b4 --- src/user/user_cache.cc | 8 ++ src/user/user_objects.cc | 153 ++++++++++++++++++++++++--------------- 2 files changed, 104 insertions(+), 57 deletions(-) diff --git a/src/user/user_cache.cc b/src/user/user_cache.cc index 64f125fc..6c97060f 100644 --- a/src/user/user_cache.cc +++ b/src/user/user_cache.cc @@ -50,6 +50,8 @@ template std::size_t mjCAsset::Add(const std::string& name, const uint8_t* data, std::size_t n); template std::size_t mjCAsset::Add(const std::string& name, const int* data, std::size_t n); +template std::size_t mjCAsset::Add(const std::string& name, + const unsigned int* data, std::size_t n); template std::size_t mjCAsset::Add(const std::string& name, const float* data, std::size_t n); template std::size_t mjCAsset::Add(const std::string& name, @@ -67,6 +69,8 @@ template std::size_t mjCAsset::AddVector(const std::string& name, const std::vector& v); template std::size_t mjCAsset::AddVector(const std::string& name, const std::vector& v); +template std::size_t mjCAsset::AddVector(const std::string& name, + const std::vector& v); template std::size_t mjCAsset::AddVector(const std::string& name, const std::vector& v); template std::size_t mjCAsset::AddVector(const std::string& name, @@ -101,6 +105,8 @@ const uint8_t* mjCAsset::Get(const std::string& name, std::size_t* n) const; template const int* mjCAsset::Get(const std::string& name, std::size_t* n) const; template +const unsigned int* mjCAsset::Get(const std::string& name, std::size_t* n) const; +template const float* mjCAsset::Get(const std::string& name, std::size_t* n) const; template const double* mjCAsset::Get(const std::string& name, std::size_t* n) const; @@ -122,6 +128,8 @@ template std::optional> mjCAsset::GetVector(const std::string& name) const; template std::optional> mjCAsset::GetVector(const std::string& name) const; +template std::optional> +mjCAsset::GetVector(const std::string& name) const; template std::optional> mjCAsset::GetVector(const std::string& name) const; template std::optional> diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 636b4921..18c5c61d 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -47,6 +47,7 @@ #include "engine/engine_util_spatial.h" #include "engine/engine_vfs.h" #include "user/user_api.h" +#include "user/user_cache.h" #include "user/user_model.h" #include "user/user_util.h" @@ -133,6 +134,91 @@ const char* FullInertia(double quat[4], double inertia[3], const double fulliner } + +// fetches cached image from PNG asset, returns nullopt if not available +static std::optional> +LoadCachedPNG(mjCAsset& asset, unsigned& w, unsigned& h, LodePNGColorType color_type) { + if (!asset.HasData("dims") || !asset.HasData("image")) { + return std::nullopt; + } + + std::vector dims + = asset.GetVector("dims").value(); + auto maybe_cached_image = asset.GetVector("image"); + + if (dims.size() != 3) { + return std::nullopt; + } + + if ((maybe_cached_image->size() != dims[0] * dims[1]) || (color_type != dims[2])) { + return std::nullopt; + } + + w = dims[0]; + h = dims[1]; + + return maybe_cached_image; +} + + + +// decodes PNG images from the given resource +std::vector LoadPNG(const mjCBase* obj, mjResource* resource, + unsigned& w, unsigned& h, LodePNGColorType color_type) { + mjCCache *cache = reinterpret_cast(mj_globalCache()); + + // try loading from cache + if (cache) { + auto asset = cache->Get(resource->name); + if (asset.has_value() && + !mju_isModifiedResource(resource, asset->Timestamp().c_str())) { + auto maybe_cached_image = LoadCachedPNG(asset.value(), w, h, color_type); + if (maybe_cached_image.has_value()) { + return maybe_cached_image.value(); + } + } + } + + // open PNG resource + const unsigned char* buffer; + int buffer_sz = mju_readResource(resource, (const void**) &buffer); + + if (buffer_sz < 0) { + throw mjCError(obj, "could not read PNG file '%s'", resource->name); + } + + if (!buffer_sz) { + throw mjCError(obj, "empty PNG file '%s'", resource->name); + } + + // decode PNG from buffer + std::vector image; + unsigned err = lodepng::decode(image, w, h, buffer, buffer_sz, color_type, 8); + + // check for errors + if (err) { + std::stringstream ss; + ss << "error decoding PNG file '" << resource->name << "': " << lodepng_error_text(err); + throw mjCError(obj, "%s", ss.str().c_str()); + } + + if (!w || !h) { + throw mjCError(obj, "error decoding PNG file '%s': zero dimension", resource->name); + } + + // insert raw image data into cache + if (cache) { + mjCAsset asset("", resource->name, resource->timestamp); + asset.AddVector("dims", std::vector{w, h, static_cast(color_type)}); + asset.AddVector("image", image); + cache->Insert(std::move(asset)); + } + + return image; +} + + + //------------------------- class mjCError implementation ------------------------------------------ // constructor @@ -2938,45 +3024,19 @@ void mjCHField::LoadCustom(mjResource* resource) { // load elevation data from PNG format void mjCHField::LoadPNG(mjResource* resource) { - // determine data source - const void* inbuffer = 0; - int inbuffer_sz = mju_readResource(resource, &inbuffer); + unsigned w, h; + std::vector image = ::LoadPNG(this, resource, w, h, LCT_GREY); - if (inbuffer_sz < 1) { - throw mjCError(this, "could not read hfield PNG file '%s'", resource->name); - } - - if (!inbuffer_sz) { - throw mjCError(this, "empty hfield PNG file '%s'", resource->name); - } - - // load PNG from file or memory - unsigned int w, h, err; - std::vector image; - err = lodepng::decode(image, w, h, (const unsigned char*) inbuffer, inbuffer_sz, LCT_GREY, 8); - - // check - if (err) { - throw mjCError(this, "PNG load error '%s' in hfield id = %d", lodepng_error_text(err), id); - } - if (!w || !h) { - throw mjCError(this, "Zero dimension in PNG hfield '%s' (id = %d)", resource->name, id); - } - - // allocate - data.assign(w*h, 0); - if (data.empty()) { - throw mjCError(this, "could not allocate buffers in hfield"); - } - - // assign and copy ncol = w; nrow = h; - for (int c=0; c& image, unsigned int& w, unsigned int& h) { - const void* inbuffer = 0; - int inbuffer_sz = mju_readResource(resource, &inbuffer); - - // still not found - if (inbuffer_sz < 1) { - throw mjCError(this, "could not read PNG texture file '%s'", resource->name); - } else if (!inbuffer_sz) { - throw mjCError(this, "PNG texture file '%s' is empty", resource->name); - } - - - // load PNG from file or memory - unsigned int err = lodepng::decode(image, w, h, (const unsigned char*) inbuffer, inbuffer_sz, LCT_RGB, 8); - - // check - if (err) { - throw mjCError(this, - "PNG file load error '%s' in texture id = %d", lodepng_error_text(err), id); - } - if (w<1 || h<1) { - throw mjCError(this, "Empty PNG file in texture '%s' (id %d)", resource->name, id); - } + image = ::LoadPNG(this, resource, w, h, LCT_RGB); }