diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD index b94623ca..36527736 100644 --- a/mediapipe/gpu/BUILD +++ b/mediapipe/gpu/BUILD @@ -368,6 +368,7 @@ cc_library( ], }), deps = [ + ":cv_texture_cache_manager", ":gpu_buffer_format", ":multi_pool", ":pixel_buffer_pool_util", diff --git a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc index b1c135af..d8155f5c 100644 --- a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc +++ b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc @@ -25,13 +25,13 @@ namespace mediapipe { CvPixelBufferPoolWrapper::CvPixelBufferPoolWrapper( int width, int height, GpuBufferFormat format, CFTimeInterval maxAge, - std::function flush_texture_caches) { + CvTextureCacheManager* texture_caches) { OSType cv_format = CVPixelFormatForGpuBufferFormat(format); CHECK_NE(cv_format, -1) << "unsupported pixel format"; pool_ = MakeCFHolderAdopting( /* keep count is 0 because the age param keeps buffers around anyway */ CreateCVPixelBufferPool(width, height, cv_format, 0, maxAge)); - flush_texture_caches_ = std::move(flush_texture_caches); + texture_caches_ = texture_caches; } CFHolder CvPixelBufferPoolWrapper::GetBuffer() { @@ -47,12 +47,12 @@ CFHolder CvPixelBufferPoolWrapper::GetBuffer() { kCFAllocatorDefault, *pool_, (__bridge CFDictionaryRef)auxAttributes, &buffer); if (err != kCVReturnWouldExceedAllocationThreshold) break; - if (flush_texture_caches_ && !tried_flushing) { + if (texture_caches_ && !tried_flushing) { // Call the flush function to potentially release old holds on buffers // and try again to create a pixel buffer. // This is used to flush CV texture caches, which may retain buffers until // flushed. - flush_texture_caches_(); + texture_caches_->FlushTextureCaches(); tried_flushing = true; } else { ++threshold; diff --git a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h index 185ba37c..7d0aec4e 100644 --- a/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h +++ b/mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h @@ -23,6 +23,7 @@ #define MEDIAPIPE_GPU_CV_PIXEL_BUFFER_POOL_WRAPPER_H_ #include "CoreFoundation/CFBase.h" +#include "mediapipe/gpu/cv_texture_cache_manager.h" #include "mediapipe/gpu/gpu_buffer_format.h" #include "mediapipe/gpu/multi_pool.h" #include "mediapipe/gpu/pixel_buffer_pool_util.h" @@ -34,15 +35,14 @@ class CvPixelBufferPoolWrapper { public: CvPixelBufferPoolWrapper(int width, int height, GpuBufferFormat format, CFTimeInterval maxAge, - std::function flush_texture_caches); + CvTextureCacheManager* texture_caches); static std::shared_ptr Create( int width, int height, GpuBufferFormat format, const MultiPoolOptions& options, - std::function flush_texture_caches = nullptr) { + CvTextureCacheManager* texture_caches = nullptr) { return std::make_shared( - width, height, format, options.max_inactive_buffer_age, - flush_texture_caches); + width, height, format, options.max_inactive_buffer_age, texture_caches); } CFHolder GetBuffer(); @@ -58,7 +58,7 @@ class CvPixelBufferPoolWrapper { private: CFHolder pool_; int count_ = 0; - std::function flush_texture_caches_; + CvTextureCacheManager* texture_caches_; }; } // namespace mediapipe diff --git a/mediapipe/gpu/gpu_buffer_multi_pool.cc b/mediapipe/gpu/gpu_buffer_multi_pool.cc index df228b7d..744ccea2 100644 --- a/mediapipe/gpu/gpu_buffer_multi_pool.cc +++ b/mediapipe/gpu/gpu_buffer_multi_pool.cc @@ -24,15 +24,10 @@ namespace mediapipe { std::shared_ptr -GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec, - const MultiPoolOptions& options) { -#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER - return CvPixelBufferPoolWrapper::Create(spec.width, spec.height, spec.format, - options, flush_platform_caches_); -#else - return GlTextureBufferPool::Create(spec.width, spec.height, spec.format, - options); -#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER +GpuBufferMultiPool::DefaultMakeSimplePool( + const GpuBufferMultiPool::BufferSpec& spec, + const MultiPoolOptions& options) { + return SimplePool::Create(spec.width, spec.height, spec.format, options); } std::shared_ptr GpuBufferMultiPool::RequestPool( @@ -44,7 +39,7 @@ std::shared_ptr GpuBufferMultiPool::RequestPool( pool = cache_.Lookup(spec, [this](const BufferSpec& spec, int request_count) { return (request_count >= options_.min_requests_before_pool) - ? MakeSimplePool(spec, options_) + ? create_simple_pool_(spec, options_) : nullptr; }); evicted = cache_.Evict(options_.max_pool_count, diff --git a/mediapipe/gpu/gpu_buffer_multi_pool.h b/mediapipe/gpu/gpu_buffer_multi_pool.h index 3ea299f7..88428d05 100644 --- a/mediapipe/gpu/gpu_buffer_multi_pool.h +++ b/mediapipe/gpu/gpu_buffer_multi_pool.h @@ -52,10 +52,6 @@ class GpuBufferMultiPool { GpuBuffer GetBuffer(int width, int height, GpuBufferFormat format = GpuBufferFormat::kBGRA32); - void SetFlushPlatformCaches(std::function flush_platform_caches) { - flush_platform_caches_ = flush_platform_caches; - } - // This class is not intended as part of the public api of this class. It is // public only because it is used as a map key type, and the map // implementation needs access to, e.g., the equality operator. @@ -74,14 +70,21 @@ class GpuBufferMultiPool { mediapipe::GpuBufferFormat format; }; - private: #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER using SimplePool = CvPixelBufferPoolWrapper; #else using SimplePool = GlTextureBufferPool; #endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER - std::shared_ptr MakeSimplePool( + using SimplePoolFactory = std::function( + const BufferSpec& spec, const MultiPoolOptions& options)>; + + void SetSimplePoolFactory(SimplePoolFactory create_simple_pool) { + create_simple_pool_ = create_simple_pool; + } + + private: + static std::shared_ptr DefaultMakeSimplePool( const GpuBufferMultiPool::BufferSpec& spec, const MultiPoolOptions& options); @@ -94,8 +97,7 @@ class GpuBufferMultiPool { absl::Mutex mutex_; mediapipe::ResourceCache> cache_ ABSL_GUARDED_BY(mutex_); - // This is used to hook up the TextureCacheManager on Apple platforms. - std::function flush_platform_caches_; + SimplePoolFactory create_simple_pool_ = DefaultMakeSimplePool; }; // BufferSpec equality operators diff --git a/mediapipe/gpu/gpu_shared_data_internal.cc b/mediapipe/gpu/gpu_shared_data_internal.cc index 457b04fd..6633c2f0 100644 --- a/mediapipe/gpu/gpu_shared_data_internal.cc +++ b/mediapipe/gpu/gpu_shared_data_internal.cc @@ -86,8 +86,12 @@ GpuResources::GpuResources(std::shared_ptr gl_context) { std::make_shared(gl_context.get()); #if __APPLE__ texture_caches_ = std::make_shared(); - gpu_buffer_pool().SetFlushPlatformCaches( - [tc = texture_caches_] { tc->FlushTextureCaches(); }); + gpu_buffer_pool().SetSimplePoolFactory( + [tc = texture_caches_](const GpuBufferMultiPool::BufferSpec& spec, + const MultiPoolOptions& options) { + return CvPixelBufferPoolWrapper::Create(spec.width, spec.height, + spec.format, options, tc.get()); + }); #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER texture_caches_->RegisterTextureCache(gl_context->cv_texture_cache()); #endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER