220 lines
7.6 KiB
C++
220 lines
7.6 KiB
C++
// Copyright 2019 The MediaPipe Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "mediapipe/framework/formats/image_multi_pool.h"
|
|
|
|
#include <tuple>
|
|
|
|
#include "absl/memory/memory.h"
|
|
#include "absl/synchronization/mutex.h"
|
|
#include "mediapipe/framework/port/logging.h"
|
|
|
|
#if !MEDIAPIPE_DISABLE_GPU
|
|
#ifdef __APPLE__
|
|
#include "mediapipe/objc/CFHolder.h"
|
|
#endif // __APPLE__
|
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
|
|
|
namespace mediapipe {
|
|
|
|
// Keep this many buffers allocated for a given frame size.
|
|
static constexpr int kKeepCount = 2;
|
|
// The maximum size of the ImageMultiPool. When the limit is reached, the
|
|
// oldest IBufferSpec will be dropped.
|
|
static constexpr int kMaxPoolCount = 20;
|
|
|
|
#if !MEDIAPIPE_DISABLE_GPU
|
|
|
|
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
|
|
|
ImageMultiPool::SimplePoolGpu ImageMultiPool::MakeSimplePoolGpu(
|
|
IBufferSpec spec) {
|
|
OSType cv_format = mediapipe::CVPixelFormatForGpuBufferFormat(
|
|
GpuBufferFormatForImageFormat(spec.format));
|
|
CHECK_NE(cv_format, -1) << "unsupported pixel format";
|
|
return MakeCFHolderAdopting(mediapipe::CreateCVPixelBufferPool(
|
|
spec.width, spec.height, cv_format, kKeepCount,
|
|
0.1 /* max age in seconds */));
|
|
}
|
|
|
|
Image ImageMultiPool::GetBufferFromSimplePool(
|
|
IBufferSpec spec, const ImageMultiPool::SimplePoolGpu& pool) {
|
|
#if TARGET_IPHONE_SIMULATOR
|
|
// On the simulator, syncing the texture with the pixelbuffer does not work,
|
|
// and we have to use glReadPixels. Since GL_UNPACK_ROW_LENGTH is not
|
|
// available in OpenGL ES 2, we should create the buffer so the pixels are
|
|
// contiguous.
|
|
//
|
|
// TODO: verify if we can use kIOSurfaceBytesPerRow to force the
|
|
// pool to give us contiguous data.
|
|
OSType cv_format = mediapipe::CVPixelFormatForGpuBufferFormat(
|
|
mediapipe::GpuBufferFormatForImageFormat(spec.format));
|
|
CHECK_NE(cv_format, -1) << "unsupported pixel format";
|
|
CVPixelBufferRef buffer;
|
|
CVReturn err = mediapipe::CreateCVPixelBufferWithoutPool(
|
|
spec.width, spec.height, cv_format, &buffer);
|
|
CHECK(!err) << "Error creating pixel buffer: " << err;
|
|
return Image(MakeCFHolderAdopting(buffer));
|
|
#else
|
|
CVPixelBufferRef buffer;
|
|
// TODO: allow the keepCount and the allocation threshold to be set
|
|
// by the application, and to be set independently.
|
|
static CFDictionaryRef auxAttributes =
|
|
mediapipe::CreateCVPixelBufferPoolAuxiliaryAttributesForThreshold(
|
|
kKeepCount);
|
|
CVReturn err = mediapipe::CreateCVPixelBufferWithPool(
|
|
*pool, auxAttributes,
|
|
[this]() {
|
|
absl::MutexLock lock(&mutex_gpu_);
|
|
for (const auto& cache : texture_caches_) {
|
|
#if TARGET_OS_OSX
|
|
CVOpenGLTextureCacheFlush(*cache, 0);
|
|
#else
|
|
CVOpenGLESTextureCacheFlush(*cache, 0);
|
|
#endif // TARGET_OS_OSX
|
|
}
|
|
},
|
|
&buffer);
|
|
CHECK(!err) << "Error creating pixel buffer: " << err;
|
|
return Image(MakeCFHolderAdopting(buffer));
|
|
#endif // TARGET_IPHONE_SIMULATOR
|
|
}
|
|
|
|
#else
|
|
|
|
ImageMultiPool::SimplePoolGpu ImageMultiPool::MakeSimplePoolGpu(
|
|
IBufferSpec spec) {
|
|
return mediapipe::GlTextureBufferPool::Create(
|
|
spec.width, spec.height, GpuBufferFormatForImageFormat(spec.format),
|
|
kKeepCount);
|
|
}
|
|
|
|
Image ImageMultiPool::GetBufferFromSimplePool(
|
|
IBufferSpec spec, const ImageMultiPool::SimplePoolGpu& pool) {
|
|
return Image(pool->GetBuffer());
|
|
}
|
|
|
|
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
|
|
|
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
|
|
|
ImageMultiPool::SimplePoolCpu ImageMultiPool::MakeSimplePoolCpu(
|
|
IBufferSpec spec) {
|
|
return ImageFramePool::Create(spec.width, spec.height, spec.format,
|
|
kKeepCount);
|
|
}
|
|
|
|
Image ImageMultiPool::GetBufferFromSimplePool(
|
|
IBufferSpec spec, const ImageMultiPool::SimplePoolCpu& pool) {
|
|
return Image(pool->GetBuffer());
|
|
}
|
|
|
|
Image ImageMultiPool::GetBuffer(int width, int height, bool use_gpu,
|
|
ImageFormat::Format format) {
|
|
#if !MEDIAPIPE_DISABLE_GPU
|
|
if (use_gpu) {
|
|
absl::MutexLock lock(&mutex_gpu_);
|
|
IBufferSpec key(width, height, format);
|
|
auto pool_it = pools_gpu_.find(key);
|
|
if (pool_it == pools_gpu_.end()) {
|
|
// Discard the least recently used pool in LRU cache.
|
|
if (pools_gpu_.size() >= kMaxPoolCount) {
|
|
auto old_spec = buffer_specs_gpu_.front(); // Front has LRU.
|
|
buffer_specs_gpu_.pop_front();
|
|
pools_gpu_.erase(old_spec);
|
|
}
|
|
buffer_specs_gpu_.push_back(key); // Push new spec to back.
|
|
std::tie(pool_it, std::ignore) = pools_gpu_.emplace(
|
|
std::piecewise_construct, std::forward_as_tuple(key),
|
|
std::forward_as_tuple(MakeSimplePoolGpu(key)));
|
|
} else {
|
|
// Find and move current 'key' spec to back, keeping others in same order.
|
|
auto specs_it = buffer_specs_gpu_.begin();
|
|
while (specs_it != buffer_specs_gpu_.end()) {
|
|
if (*specs_it == key) {
|
|
buffer_specs_gpu_.erase(specs_it);
|
|
break;
|
|
}
|
|
++specs_it;
|
|
}
|
|
buffer_specs_gpu_.push_back(key);
|
|
}
|
|
return GetBufferFromSimplePool(pool_it->first, pool_it->second);
|
|
} else // NOLINT(readability/braces)
|
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
|
{
|
|
absl::MutexLock lock(&mutex_cpu_);
|
|
IBufferSpec key(width, height, format);
|
|
auto pool_it = pools_cpu_.find(key);
|
|
if (pool_it == pools_cpu_.end()) {
|
|
// Discard the least recently used pool in LRU cache.
|
|
if (pools_cpu_.size() >= kMaxPoolCount) {
|
|
auto old_spec = buffer_specs_cpu_.front(); // Front has LRU.
|
|
buffer_specs_cpu_.pop_front();
|
|
pools_cpu_.erase(old_spec);
|
|
}
|
|
buffer_specs_cpu_.push_back(key); // Push new spec to back.
|
|
std::tie(pool_it, std::ignore) = pools_cpu_.emplace(
|
|
std::piecewise_construct, std::forward_as_tuple(key),
|
|
std::forward_as_tuple(MakeSimplePoolCpu(key)));
|
|
} else {
|
|
// Find and move current 'key' spec to back, keeping others in same order.
|
|
auto specs_it = buffer_specs_cpu_.begin();
|
|
while (specs_it != buffer_specs_cpu_.end()) {
|
|
if (*specs_it == key) {
|
|
buffer_specs_cpu_.erase(specs_it);
|
|
break;
|
|
}
|
|
++specs_it;
|
|
}
|
|
buffer_specs_cpu_.push_back(key);
|
|
}
|
|
return GetBufferFromSimplePool(pool_it->first, pool_it->second);
|
|
}
|
|
}
|
|
|
|
ImageMultiPool::~ImageMultiPool() {
|
|
#if !MEDIAPIPE_DISABLE_GPU
|
|
#ifdef __APPLE__
|
|
CHECK_EQ(texture_caches_.size(), 0)
|
|
<< "Failed to unregister texture caches before deleting pool";
|
|
#endif // defined(__APPLE__)
|
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
|
}
|
|
|
|
#if !MEDIAPIPE_DISABLE_GPU
|
|
#ifdef __APPLE__
|
|
void ImageMultiPool::RegisterTextureCache(mediapipe::CVTextureCacheType cache) {
|
|
absl::MutexLock lock(&mutex_gpu_);
|
|
|
|
CHECK(std::find(texture_caches_.begin(), texture_caches_.end(), cache) ==
|
|
texture_caches_.end())
|
|
<< "Attempting to register a texture cache twice";
|
|
texture_caches_.emplace_back(cache);
|
|
}
|
|
|
|
void ImageMultiPool::UnregisterTextureCache(
|
|
mediapipe::CVTextureCacheType cache) {
|
|
absl::MutexLock lock(&mutex_gpu_);
|
|
|
|
auto it = std::find(texture_caches_.begin(), texture_caches_.end(), cache);
|
|
CHECK(it != texture_caches_.end())
|
|
<< "Attempting to unregister an unknown texture cache";
|
|
texture_caches_.erase(it);
|
|
}
|
|
#endif // defined(__APPLE__)
|
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
|
|
|
} // namespace mediapipe
|