mediapipe/mediapipe2/examples/desktop/autoflip/quality/frame_crop_region_computer.cc
2021-06-10 23:01:19 +00:00

262 lines
10 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/examples/desktop/autoflip/quality/frame_crop_region_computer.h"
#include <cmath>
#include "mediapipe/examples/desktop/autoflip/quality/utils.h"
#include "mediapipe/framework/port/ret_check.h"
namespace mediapipe {
namespace autoflip {
absl::Status FrameCropRegionComputer::ExpandSegmentUnderConstraint(
const Segment& segment_to_add, const Segment& base_segment,
const int max_length, Segment* combined_segment,
CoverType* cover_type) const {
RET_CHECK(combined_segment != nullptr) << "Combined segment is null.";
RET_CHECK(cover_type != nullptr) << "Cover type is null.";
const LeftPoint segment_to_add_left = segment_to_add.first;
const RightPoint segment_to_add_right = segment_to_add.second;
RET_CHECK(segment_to_add_right >= segment_to_add_left)
<< "Invalid segment to add.";
const LeftPoint base_segment_left = base_segment.first;
const RightPoint base_segment_right = base_segment.second;
RET_CHECK(base_segment_right >= base_segment_left) << "Invalid base segment.";
const int base_length = base_segment_right - base_segment_left;
RET_CHECK(base_length <= max_length)
<< "Base segment length exceeds max length.";
const int segment_to_add_length = segment_to_add_right - segment_to_add_left;
const int max_leftout_amount =
std::ceil((1.0 - options_.non_required_region_min_coverage_fraction()) *
segment_to_add_length / 2);
const LeftPoint min_coverage_segment_to_add_left =
segment_to_add_left + max_leftout_amount;
const LeftPoint min_coverage_segment_to_add_right =
segment_to_add_right - max_leftout_amount;
LeftPoint combined_segment_left =
std::min(segment_to_add_left, base_segment_left);
RightPoint combined_segment_right =
std::max(segment_to_add_right, base_segment_right);
LeftPoint min_coverage_combined_segment_left =
std::min(min_coverage_segment_to_add_left, base_segment_left);
RightPoint min_coverage_combined_segment_right =
std::max(min_coverage_segment_to_add_right, base_segment_right);
if ((combined_segment_right - combined_segment_left) <= max_length) {
*cover_type = FULLY_COVERED;
} else if (min_coverage_combined_segment_right -
min_coverage_combined_segment_left <=
max_length) {
*cover_type = PARTIALLY_COVERED;
combined_segment_left = min_coverage_combined_segment_left;
combined_segment_right = min_coverage_combined_segment_right;
} else {
*cover_type = NOT_COVERED;
combined_segment_left = base_segment_left;
combined_segment_right = base_segment_right;
}
*combined_segment =
std::make_pair(combined_segment_left, combined_segment_right);
return absl::OkStatus();
}
absl::Status FrameCropRegionComputer::ExpandRectUnderConstraints(
const Rect& rect_to_add, const int max_width, const int max_height,
Rect* base_rect, CoverType* cover_type) const {
RET_CHECK(base_rect != nullptr) << "Base rect is null.";
RET_CHECK(cover_type != nullptr) << "Cover type is null.";
RET_CHECK(base_rect->width() <= max_width &&
base_rect->height() <= max_height)
<< "Base rect already exceeds target size.";
const LeftPoint rect_to_add_left = rect_to_add.x();
const RightPoint rect_to_add_right = rect_to_add.x() + rect_to_add.width();
const LeftPoint rect_to_add_top = rect_to_add.y();
const RightPoint rect_to_add_bottom = rect_to_add.y() + rect_to_add.height();
const LeftPoint base_rect_left = base_rect->x();
const RightPoint base_rect_right = base_rect->x() + base_rect->width();
const LeftPoint base_rect_top = base_rect->y();
const RightPoint base_rect_bottom = base_rect->y() + base_rect->height();
Segment horizontal_combined_segment, vertical_combined_segment;
CoverType horizontal_cover_type, vertical_cover_type;
const auto horizontal_status = ExpandSegmentUnderConstraint(
std::make_pair(rect_to_add_left, rect_to_add_right),
std::make_pair(base_rect_left, base_rect_right), max_width,
&horizontal_combined_segment, &horizontal_cover_type);
MP_RETURN_IF_ERROR(horizontal_status);
const auto vertical_status = ExpandSegmentUnderConstraint(
std::make_pair(rect_to_add_top, rect_to_add_bottom),
std::make_pair(base_rect_top, base_rect_bottom), max_height,
&vertical_combined_segment, &vertical_cover_type);
MP_RETURN_IF_ERROR(vertical_status);
if (horizontal_cover_type == NOT_COVERED ||
vertical_cover_type == NOT_COVERED) {
// Gives up if the segment is not covered in either direction.
*cover_type = NOT_COVERED;
} else {
// Tries to (partially) cover the new rect to be added.
base_rect->set_x(horizontal_combined_segment.first);
base_rect->set_y(vertical_combined_segment.first);
base_rect->set_width(horizontal_combined_segment.second -
horizontal_combined_segment.first);
base_rect->set_height(vertical_combined_segment.second -
vertical_combined_segment.first);
if (horizontal_cover_type == FULLY_COVERED &&
vertical_cover_type == FULLY_COVERED) {
*cover_type = FULLY_COVERED;
} else {
*cover_type = PARTIALLY_COVERED;
}
}
return absl::OkStatus();
}
void FrameCropRegionComputer::UpdateCropRegionScore(
const KeyFrameCropOptions::ScoreAggregationType score_aggregation_type,
const float feature_score, const bool is_required,
float* crop_region_score) {
if (feature_score < 0.0) {
LOG(WARNING) << "Ignoring negative score";
return;
}
switch (score_aggregation_type) {
case KeyFrameCropOptions::MAXIMUM: {
*crop_region_score = std::max(feature_score, *crop_region_score);
break;
}
case KeyFrameCropOptions::SUM_REQUIRED: {
if (is_required) {
*crop_region_score += feature_score;
}
break;
}
case KeyFrameCropOptions::SUM_ALL: {
*crop_region_score += feature_score;
break;
}
case KeyFrameCropOptions::CONSTANT: {
*crop_region_score = 1.0;
break;
}
default: {
LOG(WARNING) << "Unknown CropRegionScoreType " << score_aggregation_type;
break;
}
}
}
absl::Status FrameCropRegionComputer::ComputeFrameCropRegion(
const KeyFrameInfo& frame_info, KeyFrameCropResult* crop_result) const {
RET_CHECK(crop_result != nullptr) << "KeyFrameCropResult is null.";
// Set timestamp of KeyFrameCropResult
crop_result->set_timestamp_ms(frame_info.timestamp_ms());
// Sorts required and non-required regions.
std::vector<SalientRegion> required_regions, non_required_regions;
const auto sort_status = SortDetections(
frame_info.detections(), &required_regions, &non_required_regions);
MP_RETURN_IF_ERROR(sort_status);
int target_width = options_.target_width();
int target_height = options_.target_height();
auto* region = crop_result->mutable_region();
bool crop_region_is_empty = true;
float crop_region_score = 0.0;
// Gets union of all required regions.
for (int i = 0; i < required_regions.size(); ++i) {
const Rect& required_region = required_regions[i].location();
if (crop_region_is_empty) {
*region = required_region;
crop_region_is_empty = false;
} else {
RectUnion(required_region, region);
}
UpdateCropRegionScore(options_.score_aggregation_type(),
required_regions[i].score(), true,
&crop_region_score);
}
crop_result->set_required_region_is_empty(crop_region_is_empty);
if (!crop_region_is_empty) {
*crop_result->mutable_required_region() = *region;
crop_result->set_are_required_regions_covered_in_target_size(
region->width() <= target_width && region->height() <= target_height);
target_width = std::max(target_width, region->width());
target_height = std::max(target_height, region->height());
} else {
crop_result->set_are_required_regions_covered_in_target_size(true);
}
// Tries to fit non-required regions.
int num_covered = 0;
for (int i = 0; i < non_required_regions.size(); ++i) {
const Rect& non_required_region = non_required_regions[i].location();
CoverType cover_type = NOT_COVERED;
if (crop_region_is_empty) {
// If the crop region is empty, tries to expand an empty base region
// at the center of this region to include itself.
region->set_x(non_required_region.x() + non_required_region.width() / 2);
region->set_y(non_required_region.y() + non_required_region.height() / 2);
region->set_width(0);
region->set_height(0);
MP_RETURN_IF_ERROR(ExpandRectUnderConstraints(non_required_region,
target_width, target_height,
region, &cover_type));
if (cover_type != NOT_COVERED) {
crop_region_is_empty = false;
}
} else {
// Otherwise tries to expand the crop region to cover the non-required
// region under target size constraint.
MP_RETURN_IF_ERROR(ExpandRectUnderConstraints(non_required_region,
target_width, target_height,
region, &cover_type));
}
// Updates number of covered non-required regions and score.
if (cover_type == FULLY_COVERED) {
num_covered++;
UpdateCropRegionScore(options_.score_aggregation_type(),
non_required_regions[i].score(), false,
&crop_region_score);
}
}
const float fraction_covered =
non_required_regions.empty()
? 0.0
: static_cast<float>(num_covered) / non_required_regions.size();
crop_result->set_fraction_non_required_covered(fraction_covered);
crop_result->set_region_is_empty(crop_region_is_empty);
crop_result->set_region_score(crop_region_score);
return absl::OkStatus();
}
} // namespace autoflip
} // namespace mediapipe