chromium/chrome/renderer/companion/visual_query/visual_query_classification_and_eligibility.cc

// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/renderer/companion/visual_query/visual_query_classification_and_eligibility.h"

#include "base/logging.h"
#include "base/memory/ptr_util.h"
#include "base/notreached.h"
#include "skia/ext/image_operations.h"
#include "third_party/skia/include/core/SkColorSpace.h"
#include "third_party/tflite/src/tensorflow/lite/kernels/builtin_op_kernels.h"
#include "third_party/tflite/src/tensorflow/lite/op_resolver.h"
#include "third_party/tflite_support/src/tensorflow_lite_support/cc/task/core/task_api_factory.h"
#include "third_party/tflite_support/src/tensorflow_lite_support/cc/task/vision/image_classifier.h"

namespace companion::visual_query {

namespace {
ImageClassifier;
constexpr char kPosition[] = …;
constexpr char kStatic[] = …;
constexpr char kZIndex[] = …;

// TODO(b/284645622): This info should be contained in the image metadata. See
// if we can get it out.
constexpr int kImageSize = …;
constexpr char kNegativeShoppingLabel[] = …;
constexpr char kNegativeSensitivityLabel[] = …;

bool OrOfThresholdingRuleLooksGood(const OrOfThresholdingRules& rules) { … }

ClassificationMetrics CalculateClassificatonMetrics(
    EligibilityModule& eligibility_module,
    const std::vector<ImageId>& first_pass_images,
    const std::vector<ImageId>& second_pass_images) { … }

// TODO(b/284645622): this really belongs in the eligibility module code.
bool EligibilitySpecLooksGood(const EligibilitySpec& eligibility_spec) { … }

std::unique_ptr<tflite::MutableOpResolver> CreateOpResolver() { … }

std::unique_ptr<ImageClassifier> CreateClassifier(std::string model_data) { … }

std::string GetModelInput(const SkBitmap& bitmap, int width, int height) { … }

// Returns the score of the first class in the classification that does not have
// the provided avoid_name.
double GetScoreOfFirstClassThatDoesNotHaveName(
    const std::string& avoid_name,
    const tflite::task::vision::Classifications& classification) { … }
}  // namespace

SingleImageGeometryFeatures
VisualClassificationAndEligibility::ExtractFeaturesForEligibility(
    const ImageId& image_identifier,
    blink::WebElement& element) { … }

std::unique_ptr<VisualClassificationAndEligibility>
VisualClassificationAndEligibility::Create(
    const std::string& model_bytes,
    const EligibilitySpec& eligibility_spec) { … }

std::vector<ImageId>
VisualClassificationAndEligibility::RunClassificationAndEligibility(
    base::flat_map<ImageId, SingleImageFeaturesAndBytes>& images,
    const gfx::SizeF& viewport_size) { … }

VisualClassificationAndEligibility::~VisualClassificationAndEligibility() =
    default;

std::pair<double, double> VisualClassificationAndEligibility::ClassifyImage(
    const SkBitmap& bitmap) { … }

VisualClassificationAndEligibility::VisualClassificationAndEligibility() =
    default;
}  // namespace companion::visual_query