Project import generated by Copybara.
PiperOrigin-RevId: 253489161
This commit is contained in:
MediaPipe Team
@@ -0,0 +1,449 @@
|
||||
// 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/util/sequence/media_sequence.h"
|
||||
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
|
||||
#include "absl/strings/str_split.h"
|
||||
#include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/util/sequence/media_sequence_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace mediasequence {
|
||||
|
||||
namespace {
|
||||
|
||||
// Decodes the image header to get metadata as strings and ints.
|
||||
bool ImageMetadata(const std::string& image_str, std::string* format_string,
|
||||
int* width, int* height, int* channels) {
|
||||
// Determine the image encoding by matching known header bytes.
|
||||
if (image_str[0] == static_cast<char>(0x89) && image_str[1] == 'P' &&
|
||||
image_str[2] == 'N' && image_str[3] == 'G') {
|
||||
*format_string = "PNG";
|
||||
} else if (image_str[0] == static_cast<char>(0xFF) &&
|
||||
image_str[1] == static_cast<char>(0xD8) &&
|
||||
image_str[image_str.size() - 2] == static_cast<char>(0xFF) &&
|
||||
image_str[image_str.size() - 1] == static_cast<char>(0xD9)) {
|
||||
*format_string = "JPEG";
|
||||
} else {
|
||||
*format_string = "UNKNOWN";
|
||||
}
|
||||
auto buf = reinterpret_cast<void*>(const_cast<char*>(image_str.data()));
|
||||
cv::Mat img = cv::imdecode(cv::Mat(/*rows=*/image_str.size(),
|
||||
/*cols=*/1, CV_8UC1, buf),
|
||||
-1 /*cv::ImreadModes::IMREAD_UNCHANGED*/);
|
||||
if (img.data == nullptr) {
|
||||
return false;
|
||||
}
|
||||
*width = img.cols;
|
||||
*height = img.rows;
|
||||
*channels = img.channels();
|
||||
return true;
|
||||
}
|
||||
|
||||
// Finds the nearest timestamp in a FeatureList of timestamps. The FeatureList
|
||||
// must contain int64 values and only the first value at each step is used.
|
||||
int NearestIndex(int64 timestamp,
|
||||
const tensorflow::FeatureList& int64_feature_list) {
|
||||
int64 closest_distance = std::numeric_limits<int64>::max();
|
||||
int index = -1;
|
||||
for (int i = 0; i < int64_feature_list.feature_size(); ++i) {
|
||||
int64 current_value = int64_feature_list.feature(i).int64_list().value(0);
|
||||
int64 current_distance = std::abs(current_value - timestamp);
|
||||
if (current_distance < closest_distance) {
|
||||
index = i;
|
||||
closest_distance = current_distance;
|
||||
}
|
||||
}
|
||||
return index;
|
||||
}
|
||||
|
||||
// Find the numerical sampling rate between two values in seconds if the input
|
||||
// timestamps are in microseconds.
|
||||
float TimestampsToRate(int64 first_timestamp, int64 second_timestamp) {
|
||||
int64 timestamp_diff = second_timestamp - first_timestamp;
|
||||
// convert from microseconds to seconds.
|
||||
float rate = 1.0 / (static_cast<float>(timestamp_diff) / 1000000);
|
||||
return rate;
|
||||
}
|
||||
|
||||
// Sets the values of "clip/number_of_frames", "clip/keyframe/index",
|
||||
// "segment/start/index" and "segment/end/index" by finding the closest
|
||||
// timestamps in the "image/timestamp" FeatureList if image timestamps are
|
||||
// present.
|
||||
::mediapipe::Status ReconcileAnnotationIndicesByImageTimestamps(
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
if (GetImageTimestampSize(*sequence) == 0) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
int index;
|
||||
|
||||
// clip/segment/index
|
||||
if (HasSegmentStartTimestamp(*sequence)) {
|
||||
int segment_size = GetSegmentStartTimestampSize(*sequence);
|
||||
RET_CHECK_EQ(GetSegmentEndTimestampSize(*sequence), segment_size)
|
||||
<< "Expected an equal number of segment timestamps, but found "
|
||||
<< "start: " << segment_size
|
||||
<< ", end: " << GetSegmentEndTimestampSize(*sequence);
|
||||
|
||||
std::vector<int64> start_indices;
|
||||
start_indices.reserve(segment_size);
|
||||
for (const int64& timestamp : GetSegmentStartTimestamp(*sequence)) {
|
||||
index = NearestIndex(timestamp,
|
||||
GetFeatureList(*sequence, kImageTimestampKey));
|
||||
start_indices.push_back(index);
|
||||
}
|
||||
SetSegmentStartIndex(start_indices, sequence);
|
||||
|
||||
std::vector<int64> end_indices;
|
||||
end_indices.reserve(segment_size);
|
||||
for (const int64& timestamp : GetSegmentEndTimestamp(*sequence)) {
|
||||
index = NearestIndex(timestamp,
|
||||
GetFeatureList(*sequence, kImageTimestampKey));
|
||||
end_indices.push_back(index);
|
||||
}
|
||||
SetSegmentEndIndex(end_indices, sequence);
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Sets the values of "image/format", "image/channels", "image/height",
|
||||
// "image/width", and "image/frame_rate" based image metadata and timestamps.
|
||||
::mediapipe::Status ReconcileMetadataImages(
|
||||
const std::string& prefix, tensorflow::SequenceExample* sequence) {
|
||||
if (GetImageEncodedSize(prefix, *sequence) == 0) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
std::string format;
|
||||
int height, width, channels;
|
||||
RET_CHECK(ImageMetadata(GetImageEncodedAt(prefix, *sequence, 0), &format,
|
||||
&width, &height, &channels))
|
||||
<< "Failure to decode image metadata of image: "
|
||||
<< GetImageEncodedAt(prefix, *sequence, 0);
|
||||
SetImageFormat(prefix, format, sequence);
|
||||
SetImageHeight(prefix, height, sequence);
|
||||
SetImageWidth(prefix, width, sequence);
|
||||
SetImageChannels(prefix, channels, sequence);
|
||||
|
||||
if (GetImageTimestampSize(prefix, *sequence) > 1) {
|
||||
float rate = TimestampsToRate(GetImageTimestampAt(prefix, *sequence, 0),
|
||||
GetImageTimestampAt(prefix, *sequence, 1));
|
||||
SetImageFrameRate(prefix, rate, sequence);
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Sets the values of "feature/${TAG}/dimensions", and
|
||||
// "feature/${TAG}/frame_rate" for each float list feature TAG. If the
|
||||
// dimensions are already present as a context feature, this method verifies
|
||||
// the number of elements in the feature. Otherwise, it will write the
|
||||
// dimensions as a 1D vector with the number of elements.
|
||||
::mediapipe::Status ReconcileMetadataFeatureFloats(
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
// Loop through all keys and see if they contain "/feature/floats"
|
||||
// If so, check dimensions and set rate.
|
||||
for (const auto& key_value : sequence->feature_lists().feature_list()) {
|
||||
const std::string& key = key_value.first;
|
||||
if (absl::StrContains(key, kFeatureFloatsKey)) {
|
||||
const auto prefix = key.substr(0, key.find(kFeatureFloatsKey) - 1);
|
||||
int number_of_elements = GetFeatureFloatsAt(prefix, *sequence, 0).size();
|
||||
if (HasFeatureDimensions(prefix, *sequence)) {
|
||||
int64 product = 1;
|
||||
for (int64 value : GetFeatureDimensions(prefix, *sequence)) {
|
||||
product *= value;
|
||||
}
|
||||
RET_CHECK_EQ(number_of_elements, product)
|
||||
<< "The number of elements in float feature_list " << prefix
|
||||
<< "/feature/floats does not match the dimensions: "
|
||||
<< number_of_elements;
|
||||
} else {
|
||||
SetFeatureDimensions(prefix, {number_of_elements}, sequence);
|
||||
}
|
||||
|
||||
if (GetFeatureTimestampSize(prefix, *sequence) > 1) {
|
||||
float rate =
|
||||
TimestampsToRate(GetFeatureTimestampAt(prefix, *sequence, 0),
|
||||
GetFeatureTimestampAt(prefix, *sequence, 1));
|
||||
SetFeatureRate(prefix, rate, sequence);
|
||||
}
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Go through all bounding box annotations and move the annotation to the
|
||||
// nearest image frame with a timestamp. If timestamps are not present, does
|
||||
// nothing. If two or more annotations are closest to the same frame, then only
|
||||
// the closest annotation is saved. This matches the behavior of downsampling
|
||||
// images streams in time.
|
||||
::mediapipe::Status ReconcileMetadataBoxAnnotations(
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
int num_bboxes = GetBBoxTimestampSize(*sequence);
|
||||
int num_frames = GetImageTimestampSize(*sequence);
|
||||
if (num_bboxes && num_frames) {
|
||||
// If no one has indicated which frames are annotated, assume annotations
|
||||
// are dense.
|
||||
if (GetBBoxIsAnnotatedSize(*sequence) == 0) {
|
||||
for (int i = 0; i < num_bboxes; ++i) {
|
||||
AddBBoxIsAnnotated(true, sequence);
|
||||
}
|
||||
}
|
||||
RET_CHECK_EQ(num_bboxes, GetBBoxIsAnnotatedSize(*sequence))
|
||||
<< "Expected number of BBox timestamps and annotation marks to match.";
|
||||
// Update num_bboxes.
|
||||
if (GetBBoxSize(*sequence) > 0) {
|
||||
auto* bbox_feature_list =
|
||||
MutableFeatureList(kRegionBBoxXMinKey, sequence);
|
||||
RET_CHECK_EQ(num_bboxes, bbox_feature_list->feature_size())
|
||||
<< "Expected number of BBox timestamps and boxes to match.";
|
||||
ClearBBoxNumRegions(sequence);
|
||||
for (int i = 0; i < num_bboxes; ++i) {
|
||||
AddBBoxNumRegions(
|
||||
bbox_feature_list->feature(i).float_list().value_size(), sequence);
|
||||
}
|
||||
}
|
||||
// Collect which timestamps currently match to which indices in timestamps.
|
||||
// skip empty timestamps.
|
||||
// Requires sorted indices.
|
||||
::std::vector<int64> box_timestamps(num_bboxes);
|
||||
int bbox_index = 0;
|
||||
for (auto& feature :
|
||||
GetFeatureList(*sequence, kRegionTimestampKey).feature()) {
|
||||
box_timestamps[bbox_index] = feature.int64_list().value(0);
|
||||
++bbox_index;
|
||||
}
|
||||
::std::vector<int32> box_is_annotated(num_bboxes);
|
||||
bbox_index = 0;
|
||||
for (auto& feature :
|
||||
GetFeatureList(*sequence, kRegionIsAnnotatedKey).feature()) {
|
||||
box_is_annotated[bbox_index] = feature.int64_list().value(0);
|
||||
++bbox_index;
|
||||
}
|
||||
::std::vector<int64> image_timestamps(num_frames);
|
||||
int frame_index = 0;
|
||||
for (auto& feature :
|
||||
GetFeatureList(*sequence, kImageTimestampKey).feature()) {
|
||||
image_timestamps[frame_index] = feature.int64_list().value(0);
|
||||
++frame_index;
|
||||
}
|
||||
// Collect which bbox timestamps are closest to which image indices.
|
||||
::std::vector<int> bbox_index_if_annotated(num_frames, -1);
|
||||
int box_index = 0;
|
||||
int image_index = 0;
|
||||
while (box_index < num_bboxes) {
|
||||
// leave unannotated boxes at -1.
|
||||
if (!box_is_annotated[box_index]) {
|
||||
box_index += 1;
|
||||
// annotated boxes should updated their closest index.
|
||||
} else if (image_index >= num_frames - 1 ||
|
||||
llabs(image_timestamps[image_index] -
|
||||
box_timestamps[box_index]) <
|
||||
llabs(image_timestamps[image_index + 1] -
|
||||
box_timestamps[box_index])) {
|
||||
// Only overwrite with a new value if no value is present or this is
|
||||
// closer in time.
|
||||
if (bbox_index_if_annotated[image_index] == -1 ||
|
||||
llabs(image_timestamps[image_index] -
|
||||
box_timestamps[bbox_index_if_annotated[image_index]]) >
|
||||
llabs(image_timestamps[image_index] -
|
||||
box_timestamps[box_index])) {
|
||||
bbox_index_if_annotated[image_index] = box_index;
|
||||
}
|
||||
box_index += 1;
|
||||
} else {
|
||||
image_index += 1;
|
||||
}
|
||||
}
|
||||
// Only update unmodified bbox timestamp if it doesn't exist to prevent
|
||||
// overwriting with modified values.
|
||||
if (!GetUnmodifiedBBoxTimestampSize(*sequence)) {
|
||||
for (int i = 0; i < num_bboxes; ++i) {
|
||||
if (GetBBoxIsAnnotatedAt(*sequence, i)) {
|
||||
AddUnmodifiedBBoxTimestamp(box_timestamps[i], sequence);
|
||||
}
|
||||
}
|
||||
}
|
||||
// store some new feature_lists in a temporary sequence
|
||||
::tensorflow::SequenceExample tmp_seq;
|
||||
for (const auto& key_value : sequence->feature_lists().feature_list()) {
|
||||
const std::string& key = key_value.first;
|
||||
if (::absl::StartsWith(key, "region/")) {
|
||||
// create a new set of values and swap them in.
|
||||
tmp_seq.Clear();
|
||||
auto* old_feature_list = MutableFeatureList(key, sequence);
|
||||
if (key != kUnmodifiedRegionTimestampKey) {
|
||||
RET_CHECK_EQ(num_bboxes, old_feature_list->feature().size())
|
||||
<< "Expected number of BBox timestamps to match number of "
|
||||
"entries "
|
||||
<< "in " << key;
|
||||
}
|
||||
auto* new_feature_list = MutableFeatureList(key, &tmp_seq);
|
||||
for (int i = 0; i < num_frames; ++i) {
|
||||
if (bbox_index_if_annotated[i] >= 0) {
|
||||
if (key == kRegionTimestampKey) {
|
||||
new_feature_list->add_feature()->mutable_int64_list()->add_value(
|
||||
image_timestamps[i]);
|
||||
} else {
|
||||
*new_feature_list->add_feature() =
|
||||
old_feature_list->feature(bbox_index_if_annotated[i]);
|
||||
}
|
||||
} else {
|
||||
// Add either a default value or an empty.
|
||||
if (key == kRegionIsAnnotatedKey) {
|
||||
new_feature_list->add_feature()->mutable_int64_list()->add_value(
|
||||
0);
|
||||
} else if (key == kRegionNumRegionsKey) {
|
||||
new_feature_list->add_feature()->mutable_int64_list()->add_value(
|
||||
0);
|
||||
} else if (key == kRegionTimestampKey) {
|
||||
new_feature_list->add_feature()->mutable_int64_list()->add_value(
|
||||
image_timestamps[i]);
|
||||
} else if (key == kUnmodifiedRegionTimestampKey) {
|
||||
// Do not add an unmodified timestamp when
|
||||
// is_annotated == false.
|
||||
} else {
|
||||
new_feature_list->add_feature(); // Adds an empty.
|
||||
}
|
||||
}
|
||||
}
|
||||
*old_feature_list = *new_feature_list;
|
||||
}
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int GetBBoxSize(const std::string& prefix,
|
||||
const tensorflow::SequenceExample& sequence) {
|
||||
return GetBBoxXMinSize(prefix, sequence);
|
||||
}
|
||||
|
||||
std::vector<::mediapipe::Location> GetBBoxAt(
|
||||
const std::string& prefix, const tensorflow::SequenceExample& sequence,
|
||||
int index) {
|
||||
std::vector<::mediapipe::Location> bboxes;
|
||||
const auto& xmins = GetBBoxXMinAt(prefix, sequence, index);
|
||||
const auto& ymins = GetBBoxYMinAt(prefix, sequence, index);
|
||||
const auto& xmaxs = GetBBoxXMaxAt(prefix, sequence, index);
|
||||
const auto& ymaxs = GetBBoxYMaxAt(prefix, sequence, index);
|
||||
for (int i = 0; i < xmins.size(); ++i) {
|
||||
bboxes.push_back(::mediapipe::Location::CreateRelativeBBoxLocation(
|
||||
xmins[i], ymins[i], xmaxs[i] - xmins[i], ymaxs[i] - ymins[i]));
|
||||
}
|
||||
return bboxes;
|
||||
}
|
||||
|
||||
void AddBBox(const std::string& prefix,
|
||||
const std::vector<::mediapipe::Location>& bboxes,
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
::std::vector<float> xmins;
|
||||
::std::vector<float> ymins;
|
||||
::std::vector<float> xmaxs;
|
||||
::std::vector<float> ymaxs;
|
||||
for (auto& bbox : bboxes) {
|
||||
const auto& rect = bbox.GetRelativeBBox();
|
||||
xmins.push_back(rect.xmin());
|
||||
ymins.push_back(rect.ymin());
|
||||
xmaxs.push_back(rect.xmax());
|
||||
ymaxs.push_back(rect.ymax());
|
||||
}
|
||||
AddBBoxXMin(prefix, xmins, sequence);
|
||||
AddBBoxYMin(prefix, ymins, sequence);
|
||||
AddBBoxXMax(prefix, xmaxs, sequence);
|
||||
AddBBoxYMax(prefix, ymaxs, sequence);
|
||||
}
|
||||
|
||||
int GetPointSize(const std::string& prefix,
|
||||
const tensorflow::SequenceExample& sequence) {
|
||||
return GetBBoxPointXSize(prefix, sequence);
|
||||
}
|
||||
|
||||
std::vector<::std::pair<float, float>> GetPointAt(
|
||||
const std::string& prefix, const tensorflow::SequenceExample& sequence,
|
||||
int index) {
|
||||
const auto& ys = GetBBoxPointYAt(prefix, sequence, index);
|
||||
const auto& xs = GetBBoxPointXAt(prefix, sequence, index);
|
||||
std::vector<::std::pair<float, float>> points(ys.size());
|
||||
for (int i = 0; i < xs.size(); ++i) {
|
||||
points[i].first = ys[i];
|
||||
points[i].second = xs[i];
|
||||
}
|
||||
return points;
|
||||
}
|
||||
|
||||
void AddPoint(const std::string& prefix,
|
||||
const std::vector<::std::pair<float, float>>& points,
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
::std::vector<float> xs;
|
||||
::std::vector<float> ys;
|
||||
for (auto& point : points) {
|
||||
ys.push_back(point.first);
|
||||
xs.push_back(point.second);
|
||||
}
|
||||
AddBBoxPointY(prefix, ys, sequence);
|
||||
AddBBoxPointX(prefix, xs, sequence);
|
||||
}
|
||||
|
||||
std::unique_ptr<mediapipe::Matrix> GetAudioFromFeatureAt(
|
||||
const std::string& prefix, const tensorflow::SequenceExample& sequence,
|
||||
int index) {
|
||||
const auto& flat_data = GetFeatureFloatsAt(prefix, sequence, index);
|
||||
CHECK(HasFeatureNumChannels(prefix, sequence))
|
||||
<< "GetAudioAt requires num_channels context to be specified as key: "
|
||||
<< merge_prefix(prefix, kFeatureNumChannelsKey);
|
||||
int num_channels = GetFeatureNumChannels(prefix, sequence);
|
||||
CHECK_EQ(flat_data.size() % num_channels, 0)
|
||||
<< "The data size is not a multiple of the number of channels: "
|
||||
<< flat_data.size() << " % " << num_channels << " = "
|
||||
<< flat_data.size() % num_channels << " for sequence index " << index;
|
||||
auto output = absl::make_unique<mediapipe::Matrix>(
|
||||
num_channels, flat_data.size() / num_channels);
|
||||
std::copy(flat_data.begin(), flat_data.end(), output->data());
|
||||
return output;
|
||||
}
|
||||
|
||||
void AddAudioAsFeature(const std::string& prefix,
|
||||
const mediapipe::Matrix& audio,
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
auto* value_list =
|
||||
MutableFeatureList(merge_prefix(prefix, kFeatureFloatsKey), sequence)
|
||||
->add_feature()
|
||||
->mutable_float_list()
|
||||
->mutable_value();
|
||||
mediapipe::proto_ns::RepeatedField<float>(
|
||||
audio.data(), audio.data() + audio.rows() * audio.cols())
|
||||
.Swap(value_list);
|
||||
}
|
||||
|
||||
::mediapipe::Status ReconcileMetadata(bool reconcile_bbox_annotations,
|
||||
tensorflow::SequenceExample* sequence) {
|
||||
RET_CHECK_OK(ReconcileAnnotationIndicesByImageTimestamps(sequence));
|
||||
RET_CHECK_OK(ReconcileMetadataImages("", sequence));
|
||||
RET_CHECK_OK(ReconcileMetadataImages(kForwardFlowPrefix, sequence));
|
||||
RET_CHECK_OK(ReconcileMetadataImages(kClassSegmentationPrefix, sequence));
|
||||
RET_CHECK_OK(ReconcileMetadataImages(kInstanceSegmentationPrefix, sequence));
|
||||
RET_CHECK_OK(ReconcileMetadataFeatureFloats(sequence));
|
||||
if (reconcile_bbox_annotations) {
|
||||
RET_CHECK_OK(ReconcileMetadataBoxAnnotations(sequence));
|
||||
}
|
||||
// audio is always reconciled in the framework.
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediasequence
|
||||
} // namespace mediapipe
|
||||
Reference in New Issue
Block a user