Files
mediapipe/mediapipe/gpu/gpu_buffer.cc
T
MediaPipe TeamandSebastian Schmidt f8af41b1eb Internal change
PiperOrigin-RevId: 477538515
2022-09-28 21:32:36 +00:00

93 lines
2.8 KiB
C++

#include "mediapipe/gpu/gpu_buffer.h"
#include <memory>
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mediapipe/framework/port/logging.h"
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
#include "mediapipe/objc/util.h"
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
namespace mediapipe {
namespace {
struct StorageTypeFormatter {
void operator()(std::string* out,
const std::shared_ptr<internal::GpuBufferStorage>& s) const {
absl::StrAppend(out, s->storage_type().name());
}
};
} // namespace
std::string GpuBuffer::DebugString() const {
return absl::StrCat("GpuBuffer[",
absl::StrJoin(storages_, ", ", StorageTypeFormatter()),
"]");
}
internal::GpuBufferStorage& GpuBuffer::GetStorageForView(
TypeId view_provider_type, bool for_writing) const {
const std::shared_ptr<internal::GpuBufferStorage>* chosen_storage = nullptr;
// First see if any current storage supports the view.
for (const auto& s : storages_) {
if (s->can_down_cast_to(view_provider_type)) {
chosen_storage = &s;
break;
}
}
// Then try to convert existing storages to one that does.
// TODO: choose best conversion.
if (!chosen_storage) {
for (const auto& s : storages_) {
auto converter = internal::GpuBufferStorageRegistry::Get()
.StorageConverterForViewProvider(view_provider_type,
s->storage_type());
if (converter) {
storages_.push_back(converter(s));
chosen_storage = &storages_.back();
}
}
}
if (for_writing) {
if (!chosen_storage) {
// Allocate a new storage supporting the requested view.
auto factory = internal::GpuBufferStorageRegistry::Get()
.StorageFactoryForViewProvider(view_provider_type);
if (factory) {
storages_ = {factory(width(), height(), format())};
chosen_storage = &storages_.back();
}
} else {
// Discard all other storages.
storages_ = {*chosen_storage};
chosen_storage = &storages_.back();
}
}
CHECK(chosen_storage) << "no view provider found for requested view "
<< view_provider_type.name() << "; storages available: "
<< absl::StrJoin(storages_, ", ",
StorageTypeFormatter());
DCHECK((*chosen_storage)->can_down_cast_to(view_provider_type));
return **chosen_storage;
}
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer) {
auto p = buffer.internal_storage<GpuBufferStorageCvPixelBuffer>();
if (p) return **p;
return nullptr;
}
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
#endif // !MEDIAPIPE_DISABLE_GPU
} // namespace mediapipe