talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 73d686c40057684f8bfaca285368bf1813f9fc26
This commit is contained in:
MediaPipe Team
2022-03-21 12:12:39 -07:00
committed by jqtang
parent e6c19885c6
commit cc6a2f7af6
266 changed files with 3658 additions and 1681 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/formats/BUILD
+6 -5
View File
@@ -332,15 +332,13 @@ cc_library(
"//mediapipe/framework:port",
"//mediapipe/framework:type_map",
"//mediapipe/framework/port:logging",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gpu_buffer_format",
] + select({
"//conditions:default": [
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gpu_buffer_format",
"//mediapipe/gpu:gl_texture_buffer",
],
"//mediapipe:ios": [
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gpu_buffer_format",
],
"//mediapipe/gpu:disable_gpu": [],
}) + select({
@@ -430,7 +428,10 @@ cc_test(
cc_library(
name = "tensor",
srcs = ["tensor.cc"],
srcs =
[
"tensor.cc",
],
hdrs = ["tensor.h"],
copts = select({
"//mediapipe:apple": [
mediapipe/framework/formats/image.cc
+6 -51
View File
@@ -16,48 +16,15 @@
#include "mediapipe/framework/type_map.h"
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gl_texture_view.h"
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
// TODO Refactor common code from GpuBufferToImageFrameCalculator
bool Image::ConvertToCpu() const {
if (!use_gpu_) return true; // Already on CPU.
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
image_frame_ = CreateImageFrameForCVPixelBuffer(GetCVPixelBufferRef());
#else
auto gl_texture = gpu_buffer_.GetGlTextureBufferSharedPtr();
if (!gl_texture->GetProducerContext()) return false;
gl_texture->GetProducerContext()->Run([this, &gl_texture]() {
gl_texture->WaitOnGpu();
const auto gpu_buf = mediapipe::GpuBuffer(GetGlTextureBufferSharedPtr());
#ifdef __ANDROID__
glBindFramebuffer(GL_FRAMEBUFFER, 0); // b/32091368
#endif
GLuint fb = 0;
glDisable(GL_DEPTH_TEST);
// TODO Re-use a shared framebuffer.
glGenFramebuffers(1, &fb);
glBindFramebuffer(GL_FRAMEBUFFER, fb);
glViewport(0, 0, gpu_buf.width(), gpu_buf.height());
glActiveTexture(GL_TEXTURE0);
glBindTexture(gl_texture->target(), gl_texture->name());
glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
gl_texture->target(), gl_texture->name(), 0);
auto frame = std::make_shared<ImageFrame>(
mediapipe::ImageFormatForGpuBufferFormat(gpu_buf.format()),
gpu_buf.width(), gpu_buf.height(),
ImageFrame::kGlDefaultAlignmentBoundary);
const auto info = GlTextureInfoForGpuBufferFormat(
gpu_buf.format(), 0, gl_texture->GetProducerContext()->GetGlVersion());
glReadPixels(0, 0, gpu_buf.width(), gpu_buf.height(), info.gl_format,
info.gl_type, frame->MutablePixelData());
glDeleteFramebuffers(1, &fb);
// Cleanup
gl_texture->DidRead(gl_texture->GetProducerContext()->CreateSyncToken());
image_frame_ = frame;
});
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
#endif // !MEDIAPIPE_DISABLE_GPU
auto view = gpu_buffer_.GetReadView<ImageFrame>();
use_gpu_ = false;
return true;
}
@@ -67,19 +34,7 @@ bool Image::ConvertToGpu() const {
#if MEDIAPIPE_DISABLE_GPU
return false;
#else
if (use_gpu_) return true; // Already on GPU.
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
auto packet = PointToForeign<ImageFrame>(image_frame_.get());
CFHolder<CVPixelBufferRef> buffer;
auto status = CreateCVPixelBufferForImageFramePacket(packet, true, &buffer);
CHECK_OK(status);
gpu_buffer_ = mediapipe::GpuBuffer(std::move(buffer));
#else
// GlCalculatorHelperImpl::MakeGlTextureBuffer (CreateSourceTexture)
auto buffer = mediapipe::GlTextureBuffer::Create(*image_frame_);
glFlush();
gpu_buffer_ = mediapipe::GpuBuffer(std::move(buffer));
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
auto view = gpu_buffer_.GetReadView<GlTextureView>(0);
use_gpu_ = true;
return true;
#endif // MEDIAPIPE_DISABLE_GPU
mediapipe/framework/formats/image.h
+34 -85
View File
@@ -21,20 +21,18 @@
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/port/logging.h"
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/gpu_buffer_format.h"
#include "mediapipe/gpu/gpu_buffer_storage_image_frame.h"
#include "mediapipe/gpu/image_frame_view.h"
#if !MEDIAPIPE_DISABLE_GPU
#if defined(__APPLE__)
#include <CoreVideo/CoreVideo.h>
#include "mediapipe/objc/CFHolder.h"
#include "mediapipe/objc/util.h"
#if !TARGET_OS_OSX // iOS, use CVPixelBuffer.
#define MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER 1
#endif // TARGET_OS_OSX
#endif // defined(__APPLE__)
#if !MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER // OSX, use GL textures.
@@ -73,15 +71,15 @@ class Image {
// Creates an Image representing the same image content as the ImageFrame
// the input shared pointer points to, and retaining shared ownership.
explicit Image(ImageFrameSharedPtr image_frame)
: image_frame_(std::move(image_frame)) {
: gpu_buffer_(std::make_shared<GpuBufferStorageImageFrame>(
std::move(image_frame))) {
use_gpu_ = false;
pixel_mutex_ = std::make_shared<absl::Mutex>();
}
// CPU getters.
const ImageFrameSharedPtr& GetImageFrameSharedPtr() const {
if (use_gpu_ == true) ConvertToCpu();
return image_frame_;
ImageFrameSharedPtr GetImageFrameSharedPtr() const {
// Write view currently because the return type does not point to const IF.
return gpu_buffer_.GetWriteView<ImageFrame>();
}
// Creates an Image representing the same image content as the input GPU
@@ -99,19 +97,18 @@ class Image {
explicit Image(mediapipe::GpuBuffer gpu_buffer) {
use_gpu_ = true;
gpu_buffer_ = gpu_buffer;
pixel_mutex_ = std::make_shared<absl::Mutex>();
}
// GPU getters.
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
CVPixelBufferRef GetCVPixelBufferRef() const {
if (use_gpu_ == false) ConvertToGpu();
return gpu_buffer_.GetCVPixelBufferRef();
return mediapipe::GetCVPixelBufferRef(gpu_buffer_);
}
#else
mediapipe::GlTextureBufferSharedPtr GetGlTextureBufferSharedPtr() const {
if (use_gpu_ == false) ConvertToGpu();
return gpu_buffer_.GetGlTextureBufferSharedPtr();
return gpu_buffer_.internal_storage<mediapipe::GlTextureBuffer>();
}
#endif // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
// Get a GPU view. Automatically uploads from CPU if needed.
@@ -128,9 +125,7 @@ class Image {
int step() const; // Row size in bytes.
bool UsesGpu() const { return use_gpu_; }
ImageFormat::Format image_format() const;
#if !MEDIAPIPE_DISABLE_GPU
mediapipe::GpuBufferFormat format() const;
#endif // !MEDIAPIPE_DISABLE_GPU
// Converts to true iff valid.
explicit operator bool() const { return operator!=(nullptr); }
@@ -147,8 +142,8 @@ class Image {
// Lock/Unlock pixel data.
// Should be used exclusively by the PixelLock helper class.
void LockPixels() const ABSL_EXCLUSIVE_LOCK_FUNCTION(pixel_mutex_);
void UnlockPixels() const ABSL_UNLOCK_FUNCTION(pixel_mutex_);
void LockPixels() const ABSL_EXCLUSIVE_LOCK_FUNCTION();
void UnlockPixels() const ABSL_UNLOCK_FUNCTION();
// Helper utility for GPU->CPU data transfer.
bool ConvertToCpu() const;
@@ -157,75 +152,32 @@ class Image {
bool ConvertToGpu() const;
private:
#if !MEDIAPIPE_DISABLE_GPU
mutable mediapipe::GpuBuffer gpu_buffer_;
#endif // !MEDIAPIPE_DISABLE_GPU
mutable ImageFrameSharedPtr image_frame_;
mutable bool use_gpu_ = false;
mutable std::shared_ptr<absl::Mutex> pixel_mutex_; // ImageFrame only.
};
inline int Image::width() const {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_)
return gpu_buffer_.width();
else
#endif // !MEDIAPIPE_DISABLE_GPU
return image_frame_->Width();
}
inline int Image::width() const { return gpu_buffer_.width(); }
inline int Image::height() const {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_)
return gpu_buffer_.height();
else
#endif // !MEDIAPIPE_DISABLE_GPU
return image_frame_->Height();
}
inline int Image::height() const { return gpu_buffer_.height(); }
inline ImageFormat::Format Image::image_format() const {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_)
return mediapipe::ImageFormatForGpuBufferFormat(gpu_buffer_.format());
else
#endif // !MEDIAPIPE_DISABLE_GPU
return image_frame_->Format();
return mediapipe::ImageFormatForGpuBufferFormat(gpu_buffer_.format());
}
#if !MEDIAPIPE_DISABLE_GPU
inline mediapipe::GpuBufferFormat Image::format() const {
if (use_gpu_)
return gpu_buffer_.format();
else
return mediapipe::GpuBufferFormatForImageFormat(image_frame_->Format());
return gpu_buffer_.format();
}
#endif // !MEDIAPIPE_DISABLE_GPU
inline bool Image::operator==(std::nullptr_t other) const {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_)
return gpu_buffer_ == other;
else
#endif // !MEDIAPIPE_DISABLE_GPU
return image_frame_ == other;
return gpu_buffer_ == other;
}
inline bool Image::operator==(const Image& other) const {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_)
return gpu_buffer_ == other.gpu_buffer_;
else
#endif // !MEDIAPIPE_DISABLE_GPU
return image_frame_ == other.image_frame_;
return gpu_buffer_ == other.gpu_buffer_;
}
inline Image& Image::operator=(std::nullptr_t other) {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_)
gpu_buffer_ = other;
else
#endif // !MEDIAPIPE_DISABLE_GPU
image_frame_ = other;
gpu_buffer_ = other;
return *this;
}
@@ -234,19 +186,14 @@ inline int Image::channels() const {
}
inline int Image::step() const {
if (use_gpu_)
return width() * channels() *
ImageFrame::ByteDepthForFormat(image_format());
else
return image_frame_->WidthStep();
return gpu_buffer_.GetReadView<ImageFrame>()->WidthStep();
}
inline void Image::LockPixels() const {
pixel_mutex_->Lock();
ConvertToCpu(); // Download data if necessary.
}
inline void Image::UnlockPixels() const { pixel_mutex_->Unlock(); }
inline void Image::UnlockPixels() const {}
// Helper class for getting access to Image CPU data,
// and handles automatically locking/unlocking CPU data access.
@@ -268,7 +215,10 @@ class PixelReadLock {
public:
explicit PixelReadLock(const Image& image) {
buffer_ = &image;
if (buffer_) buffer_->LockPixels();
if (buffer_) {
buffer_->LockPixels();
frame_ = buffer_->GetImageFrameSharedPtr();
}
}
~PixelReadLock() {
if (buffer_) buffer_->UnlockPixels();
@@ -276,10 +226,7 @@ class PixelReadLock {
PixelReadLock(const PixelReadLock&) = delete;
const uint8* Pixels() const {
if (buffer_ && !buffer_->UsesGpu()) {
ImageFrame* frame = buffer_->GetImageFrameSharedPtr().get();
if (frame) return frame->PixelData();
}
if (frame_) return frame_->PixelData();
return nullptr;
}
@@ -287,13 +234,17 @@ class PixelReadLock {
private:
const Image* buffer_ = nullptr;
std::shared_ptr<ImageFrame> frame_;
};
class PixelWriteLock {
public:
explicit PixelWriteLock(Image* image) {
buffer_ = image;
if (buffer_) buffer_->LockPixels();
if (buffer_) {
buffer_->LockPixels();
frame_ = buffer_->GetImageFrameSharedPtr();
}
}
~PixelWriteLock() {
if (buffer_) buffer_->UnlockPixels();
@@ -301,10 +252,7 @@ class PixelWriteLock {
PixelWriteLock(const PixelWriteLock&) = delete;
uint8* Pixels() {
if (buffer_ && !buffer_->UsesGpu()) {
ImageFrame* frame = buffer_->GetImageFrameSharedPtr().get();
if (frame) return frame->MutablePixelData();
}
if (frame_) return frame_->MutablePixelData();
return nullptr;
}
@@ -312,6 +260,7 @@ class PixelWriteLock {
private:
const Image* buffer_ = nullptr;
std::shared_ptr<ImageFrame> frame_;
};
} // namespace mediapipe
mediapipe/framework/formats/image_opencv.cc
+18 -5
View File
@@ -77,7 +77,16 @@ int GetMatType(const mediapipe::ImageFormat::Format format) {
namespace mediapipe {
namespace formats {
cv::Mat MatView(const mediapipe::Image* image) {
std::shared_ptr<cv::Mat> MatView(const mediapipe::Image* image) {
// Used to hold the lock through the Mat's lifetime.
struct MatWithPixelLock {
// Constructor needed because you cannot use aggregate initialization with
// std::make_shared.
MatWithPixelLock(mediapipe::Image* image) : lock(image) {}
mediapipe::PixelWriteLock lock;
cv::Mat mat;
};
const int dims = 2;
const int sizes[] = {image->height(), image->width()};
const int type =
@@ -85,18 +94,22 @@ cv::Mat MatView(const mediapipe::Image* image) {
const size_t steps[] = {static_cast<size_t>(image->step()),
static_cast<size_t>(ImageFrame::ByteDepthForFormat(
image->image_format()))};
mediapipe::PixelWriteLock dst_lock(const_cast<mediapipe::Image*>(image));
uint8* data_ptr = dst_lock.Pixels();
auto owner =
std::make_shared<MatWithPixelLock>(const_cast<mediapipe::Image*>(image));
uint8* data_ptr = owner->lock.Pixels();
CHECK(data_ptr != nullptr);
// Use Image to initialize in-place. Image still owns memory.
if (steps[0] == sizes[1] * image->channels() *
ImageFrame::ByteDepthForFormat(image->image_format())) {
// Contiguous memory optimization. See b/78570764
return cv::Mat(dims, sizes, type, data_ptr);
owner->mat = cv::Mat(dims, sizes, type, data_ptr);
} else {
// Custom width step.
return cv::Mat(dims, sizes, type, data_ptr, steps);
owner->mat = cv::Mat(dims, sizes, type, data_ptr, steps);
}
// Aliasing constructor makes a shared_ptr<Mat> which keeps the whole
// MatWithPixelLock alive.
return std::shared_ptr<cv::Mat>(owner, &owner->mat);
}
} // namespace formats
} // namespace mediapipe
mediapipe/framework/formats/image_opencv.h
+3 -1
View File
@@ -29,7 +29,9 @@ namespace formats {
// the const modifier is lost. The caller must be careful
// not to use the returned object to modify the data in a const Image,
// even though the returned data is mutable.
cv::Mat MatView(const mediapipe::Image* image);
// Note: this returns a shared_ptr so it can keep the CPU memory referenced
// by the Mat alive.
std::shared_ptr<cv::Mat> MatView(const mediapipe::Image* image);
} // namespace formats
} // namespace mediapipe
mediapipe/framework/formats/matrix.h
+1 -1
View File
@@ -39,7 +39,7 @@ void MatrixDataProtoFromMatrix(const Matrix& matrix, MatrixData* matrix_data);
void MatrixFromMatrixDataProto(const MatrixData& matrix_data, Matrix* matrix);
#if !defined(MEDIAPIPE_MOBILE) && !defined(MEDIAPIPE_LITE)
// Produce a Text format MatrixData std::string. Mainly useful for test code.
// Produce a Text format MatrixData string. Mainly useful for test code.
std::string MatrixAsTextProto(const Matrix& matrix);
// Produce a Matrix from a text format MatrixData proto representation.
void MatrixFromTextProto(const std::string& text_proto, Matrix* matrix);
mediapipe/framework/formats/tensor.h
+8 -1
View File
@@ -76,7 +76,7 @@ class Tensor {
public:
// No resources are allocated here.
enum class ElementType { kNone, kFloat16, kFloat32 };
enum class ElementType { kNone, kFloat16, kFloat32, kUInt8 };
struct Shape {
Shape() = default;
Shape(std::initializer_list<int> dimensions) : dims(dimensions) {}
@@ -215,6 +215,8 @@ class Tensor {
return 2;
case ElementType::kFloat32:
return sizeof(float);
case ElementType::kUInt8:
return 1;
}
}
int bytes() const { return shape_.num_elements() * element_size(); }
@@ -278,6 +280,11 @@ class Tensor {
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
};
int BhwcBatchFromShape(const Tensor::Shape& shape);
int BhwcHeightFromShape(const Tensor::Shape& shape);
int BhwcWidthFromShape(const Tensor::Shape& shape);
int BhwcDepthFromShape(const Tensor::Shape& shape);
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_H_
mediapipe/framework/formats/tensor_internal.h
+25 -2
View File
@@ -16,14 +16,22 @@
#define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_INTERNAL_H_
#include <cstdint>
#include <type_traits>
#include "mediapipe/framework/tool/type_util.h"
namespace mediapipe {
// Generates unique view id at compile-time using FILE and LINE.
#define TENSOR_UNIQUE_VIEW_TYPE_ID() \
static constexpr uint64_t kId = tensor_internal::FnvHash64( \
#define TENSOR_UNIQUE_VIEW_TYPE_ID() \
static inline uint64_t kId = tensor_internal::FnvHash64( \
__FILE__, tensor_internal::FnvHash64(TENSOR_INT_TO_STRING(__LINE__)))
// Generates unique view id at compile-time using FILE and LINE and Type of the
// template view's argument.
#define TENSOR_UNIQUE_VIEW_TYPE_ID_T(T) \
static inline uint64_t kId = tool::GetTypeHash<T>();
namespace tensor_internal {
#define TENSOR_INT_TO_STRING2(x) #x
@@ -36,6 +44,21 @@ constexpr uint64_t kFnvOffsetBias = 0xcbf29ce484222325;
constexpr uint64_t FnvHash64(const char* str, uint64_t hash = kFnvOffsetBias) {
return (str[0] == 0) ? hash : FnvHash64(str + 1, (hash ^ str[0]) * kFnvPrime);
}
template <typename... Ts>
struct TypeList {
static constexpr std::size_t size{sizeof...(Ts)};
};
template <typename, typename>
struct TypeInList {};
template <typename T, typename... Ts>
struct TypeInList<T, TypeList<T, Ts...>>
: std::integral_constant<std::size_t, 0> {};
template <typename T, typename TOther, typename... Ts>
struct TypeInList<T, TypeList<TOther, Ts...>>
: std::integral_constant<std::size_t,
1 + TypeInList<T, TypeList<Ts...>>::value> {};
} // namespace tensor_internal
} // namespace mediapipe