screen_ai_service_impl.cc 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. // Copyright 2022 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "components/services/screen_ai/screen_ai_service_impl.h"
  5. #include "base/command_line.h"
  6. #include "base/files/file_path.h"
  7. #include "base/process/process.h"
  8. #include "components/services/screen_ai/proto/proto_convertor.h"
  9. #include "components/services/screen_ai/public/cpp/utilities.h"
  10. #include "components/services/screen_ai/screen_ai_ax_tree_serializer.h"
  11. #include "sandbox/policy/switches.h"
  12. #include "ui/accessibility/accessibility_features.h"
  13. #include "ui/accessibility/ax_tree_id.h"
  14. #include "ui/gfx/geometry/rect_f.h"
  15. namespace screen_ai {
  16. namespace {
  17. base::FilePath GetLibraryFilePath() {
  18. base::FilePath library_path = GetStoredComponentBinaryPath();
  19. if (!library_path.empty())
  20. return library_path;
  21. // Binary file path is set while setting the sandbox on Linux, or the first
  22. // time this function is called. So in all other cases we need to look for
  23. // the library binary in its component folder.
  24. library_path = GetLatestComponentBinaryPath();
  25. StoreComponentBinaryPath(library_path);
  26. return library_path;
  27. }
  28. } // namespace
  29. ScreenAIService::ScreenAIService(
  30. mojo::PendingReceiver<mojom::ScreenAIService> receiver)
  31. : library_(GetLibraryFilePath()),
  32. init_function_(
  33. reinterpret_cast<InitFunction>(library_.GetFunctionPointer("Init"))),
  34. annotate_function_(reinterpret_cast<AnnotateFunction>(
  35. library_.GetFunctionPointer("Annotate"))),
  36. extract_main_content_function_(
  37. reinterpret_cast<ExtractMainContentFunction>(
  38. library_.GetFunctionPointer("ExtractMainContent"))),
  39. receiver_(this, std::move(receiver)) {
  40. DCHECK(init_function_ && annotate_function_ &&
  41. extract_main_content_function_);
  42. if (!init_function_(
  43. /*init_visual_annotations = */ features::
  44. IsScreenAIVisualAnnotationsEnabled() ||
  45. features::IsPdfOcrEnabled(),
  46. /*init_main_content_extraction = */
  47. features::IsReadAnythingWithScreen2xEnabled(),
  48. /*debug_mode = */ features::IsScreenAIDebugModeEnabled(),
  49. /*models_path = */
  50. GetLibraryFilePath().DirName().MaybeAsASCII().c_str())) {
  51. // TODO(https://crbug.com/1278249): Add UMA metrics to monitor failures.
  52. VLOG(0) << "Screen AI library initialization failed.";
  53. base::Process::TerminateCurrentProcessImmediately(-1);
  54. }
  55. }
  56. ScreenAIService::~ScreenAIService() = default;
  57. void ScreenAIService::BindAnnotator(
  58. mojo::PendingReceiver<mojom::ScreenAIAnnotator> annotator) {
  59. screen_ai_annotators_.Add(this, std::move(annotator));
  60. }
  61. void ScreenAIService::BindMainContentExtractor(
  62. mojo::PendingReceiver<mojom::Screen2xMainContentExtractor>
  63. main_content_extractor) {
  64. screen_2x_main_content_extractors_.Add(this,
  65. std::move(main_content_extractor));
  66. }
  67. void ScreenAIService::Annotate(const SkBitmap& image,
  68. AnnotationCallback callback) {
  69. ui::AXTreeUpdate update;
  70. VLOG(2) << "Screen AI library starting to process " << image.width() << "x"
  71. << image.height() << " snapshot.";
  72. char* annotation_proto = nullptr;
  73. uint32_t annotation_proto_length = 0;
  74. // TODO(https://crbug.com/1278249): Consider adding a signature that
  75. // verifies the data integrity and source.
  76. if (annotate_function_(image, annotation_proto, annotation_proto_length)) {
  77. DCHECK(annotation_proto);
  78. std::string proto_as_string;
  79. proto_as_string.resize(annotation_proto_length);
  80. memcpy(static_cast<void*>(proto_as_string.data()), annotation_proto,
  81. annotation_proto_length);
  82. delete annotation_proto;
  83. gfx::Rect image_rect(image.width(), image.height());
  84. update =
  85. ScreenAIVisualAnnotationToAXTreeUpdate(proto_as_string, image_rect);
  86. // TODO(nektar): Get the parent tree ID from the calling process (i.e.
  87. // browser or renderer).
  88. ScreenAIAXTreeSerializer serializer(
  89. /* parent_tree_id */ ui::AXTreeID::CreateNewAXTreeID(),
  90. std::move(update.nodes));
  91. update = serializer.Serialize();
  92. } else {
  93. VLOG(1) << "Screen AI library could not process snapshot.";
  94. }
  95. std::move(callback).Run(update);
  96. }
  97. void ScreenAIService::ExtractMainContent(const ui::AXTreeUpdate& snapshot,
  98. ContentExtractionCallback callback) {
  99. std::string serialized_snapshot = Screen2xSnapshotToViewHierarchy(snapshot);
  100. std::vector<int32_t> content_node_ids;
  101. int32_t* node_ids = nullptr;
  102. uint32_t nodes_count = 0;
  103. if (!extract_main_content_function_(serialized_snapshot.c_str(),
  104. serialized_snapshot.length(), node_ids,
  105. nodes_count)) {
  106. VLOG(1) << "Screen2x did not return main content.";
  107. } else {
  108. content_node_ids.resize(nodes_count);
  109. memcpy(content_node_ids.data(), node_ids, nodes_count * sizeof(int32_t));
  110. delete[] node_ids;
  111. node_ids = nullptr;
  112. VLOG(2) << "Screen2x returned " << content_node_ids.size() << " node ids:";
  113. for (int32_t i : content_node_ids)
  114. VLOG(2) << i;
  115. }
  116. std::move(callback).Run(content_node_ids);
  117. }
  118. } // namespace screen_ai