Add AV1 I444 support to CRD host

This CL adds the code needed to the CRD host to allow it provide frames
to the AV1 encoder using the I444 format.

Bug: 1329660
Change-Id: I3795a1aab734fdfe26bfc135ffc8b723a9eb6c41
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3867799
Reviewed-by: Lambros Lambrou <lambroslambrou@chromium.org>
Commit-Queue: Joe Downing <joedow@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1041829}
diff --git a/remoting/codec/webrtc_video_encoder_av1.cc b/remoting/codec/webrtc_video_encoder_av1.cc
index fa3df8c..779d99bf 100644
--- a/remoting/codec/webrtc_video_encoder_av1.cc
+++ b/remoting/codec/webrtc_video_encoder_av1.cc
@@ -10,6 +10,7 @@
 #include "base/system/sys_info.h"
 #include "remoting/base/cpu_utils.h"
 #include "remoting/base/util.h"
+#include "third_party/libaom/source/libaom/aom/aom_image.h"
 #include "third_party/libaom/source/libaom/aom/aomcx.h"
 #include "third_party/libyuv/include/libyuv/convert_from_argb.h"
 #include "third_party/webrtc/modules/desktop_capture/desktop_frame.h"
@@ -48,7 +49,10 @@
 }
 
 void WebrtcVideoEncoderAV1::SetLosslessColor(bool want_lossless) {
-  NOTIMPLEMENTED();
+  if (want_lossless != lossless_color_) {
+    lossless_color_ = want_lossless;
+    codec_.reset();
+  }
 }
 
 bool WebrtcVideoEncoderAV1::InitializeCodec(const webrtc::DesktopSize& size) {
@@ -63,6 +67,9 @@
   // values based on the frame dimensions later on.
   config_.g_threads = GetEncoderThreadCount(std::max(config_.g_w, config_.g_h));
 
+  // Choose a profile based on whether we should provide frames in I420 or I444.
+  config_.g_profile = lossless_color_ ? 1 : 0;
+
   // Initialize an encoder instance.
   scoped_aom_codec codec(new aom_codec_ctx_t, DestroyAomCodecContext);
   codec->name = nullptr;
@@ -198,7 +205,8 @@
     updated_region.IntersectWith(
         webrtc::DesktopRect::MakeWH(image_->d_w, image_->d_h));
   } else {
-    image_.reset(aom_img_alloc(nullptr, AOM_IMG_FMT_I420, frame->size().width(),
+    aom_img_fmt_t fmt = lossless_color_ ? AOM_IMG_FMT_I444 : AOM_IMG_FMT_I420;
+    image_.reset(aom_img_alloc(nullptr, fmt, frame->size().width(),
                                frame->size().height(),
                                GetSimdMemoryAlignment()));
     updated_region.AddRect(
@@ -215,18 +223,37 @@
   uint8_t* u_data = image_->planes[1];
   uint8_t* v_data = image_->planes[2];
 
-  CHECK_EQ(image_->fmt, AOM_IMG_FMT_I420);
-  for (webrtc::DesktopRegion::Iterator r(updated_region); !r.IsAtEnd();
-       r.Advance()) {
-    webrtc::DesktopRect rect = GetRowAlignedRect(r.rect(), image_->d_w);
-    int rgb_offset = rgb_stride * rect.top() +
-                     rect.left() * webrtc::DesktopFrame::kBytesPerPixel;
-    int y_offset = y_stride * rect.top() + rect.left();
-    int uv_offset = uv_stride * rect.top() / 2 + rect.left() / 2;
-    libyuv::ARGBToI420(rgb_data + rgb_offset, rgb_stride, y_data + y_offset,
-                       y_stride, u_data + uv_offset, uv_stride,
-                       v_data + uv_offset, uv_stride, rect.width(),
-                       rect.height());
+  switch (image_->fmt) {
+    case AOM_IMG_FMT_I420:
+      for (webrtc::DesktopRegion::Iterator r(updated_region); !r.IsAtEnd();
+           r.Advance()) {
+        webrtc::DesktopRect rect = GetRowAlignedRect(r.rect(), image_->d_w);
+        int rgb_offset = rgb_stride * rect.top() +
+                         rect.left() * webrtc::DesktopFrame::kBytesPerPixel;
+        int y_offset = y_stride * rect.top() + rect.left();
+        int uv_offset = uv_stride * rect.top() / 2 + rect.left() / 2;
+        libyuv::ARGBToI420(rgb_data + rgb_offset, rgb_stride, y_data + y_offset,
+                           y_stride, u_data + uv_offset, uv_stride,
+                           v_data + uv_offset, uv_stride, rect.width(),
+                           rect.height());
+      }
+      break;
+    case AOM_IMG_FMT_I444:
+      for (webrtc::DesktopRegion::Iterator r(updated_region); !r.IsAtEnd();
+           r.Advance()) {
+        webrtc::DesktopRect rect = GetRowAlignedRect(r.rect(), image_->d_w);
+        int rgb_offset = rgb_stride * rect.top() +
+                         rect.left() * webrtc::DesktopFrame::kBytesPerPixel;
+        int yuv_offset = uv_stride * rect.top() + rect.left();
+        libyuv::ARGBToI444(rgb_data + rgb_offset, rgb_stride,
+                           y_data + yuv_offset, y_stride, u_data + yuv_offset,
+                           uv_stride, v_data + yuv_offset, uv_stride,
+                           rect.width(), rect.height());
+      }
+      break;
+    default:
+      NOTREACHED();
+      break;
   }
 }
 
@@ -240,6 +267,8 @@
 
   // Encoder config values are defined in:
   // //third_party/libaom/source/libaom/aom/aom_encoder.h
+
+  // Default to profile 0 and update later if lossless color is requested.
   config_.g_profile = 0;
   // Width, height, and thread count are set once the frame size is known.
   config_.g_w = 0;
diff --git a/remoting/codec/webrtc_video_encoder_av1.h b/remoting/codec/webrtc_video_encoder_av1.h
index b4b454fa..bdb9dad 100644
--- a/remoting/codec/webrtc_video_encoder_av1.h
+++ b/remoting/codec/webrtc_video_encoder_av1.h
@@ -53,6 +53,10 @@
   using scoped_aom_image = std::unique_ptr<aom_image_t, void (*)(aom_image_t*)>;
   scoped_aom_image image_;
 
+  // Indicates whether the frames provided to the encoder will use I420 (lossy)
+  // or I444 (lossless) format.
+  bool lossless_color_ = false;
+
   // Active map used to optimize out processing of unchanged macroblocks.
   VideoEncoderActiveMap active_map_;
   // Disable |active_map_| until we've verified it improves performance.
diff --git a/remoting/protocol/webrtc_video_encoder_factory.cc b/remoting/protocol/webrtc_video_encoder_factory.cc
index e27f10a..e7a032fe 100644
--- a/remoting/protocol/webrtc_video_encoder_factory.cc
+++ b/remoting/protocol/webrtc_video_encoder_factory.cc
@@ -9,6 +9,7 @@
 #include "base/threading/thread_task_runner_handle.h"
 #include "remoting/protocol/video_channel_state_observer.h"
 #include "remoting/protocol/webrtc_video_encoder_wrapper.h"
+#include "third_party/webrtc/api/video_codecs/av1_profile.h"
 #include "third_party/webrtc/api/video_codecs/sdp_video_format.h"
 #include "third_party/webrtc/api/video_codecs/video_codec.h"
 #include "third_party/webrtc/api/video_codecs/vp9_profile.h"
@@ -35,6 +36,10 @@
   }
 #endif
   formats_.emplace_back(webrtc::SdpVideoFormat("AV1"));
+  formats_.emplace_back(webrtc::SdpVideoFormat(
+      "AV1",
+      {{webrtc::kAV1FmtpProfile,
+        webrtc::AV1ProfileToString(webrtc::AV1Profile::kProfile1).data()}}));
 }
 
 WebrtcVideoEncoderFactory::~WebrtcVideoEncoderFactory() = default;
diff --git a/remoting/protocol/webrtc_video_encoder_wrapper.cc b/remoting/protocol/webrtc_video_encoder_wrapper.cc
index e50cd14..af0f984 100644
--- a/remoting/protocol/webrtc_video_encoder_wrapper.cc
+++ b/remoting/protocol/webrtc_video_encoder_wrapper.cc
@@ -25,6 +25,7 @@
 #include "remoting/codec/webrtc_video_encoder_vpx.h"
 #include "remoting/protocol/video_channel_state_observer.h"
 #include "remoting/protocol/webrtc_video_frame_adapter.h"
+#include "third_party/webrtc/api/video_codecs/av1_profile.h"
 #include "third_party/webrtc/api/video_codecs/sdp_video_format.h"
 #include "third_party/webrtc/api/video_codecs/vp9_profile.h"
 #include "third_party/webrtc/modules/desktop_capture/desktop_frame.h"
@@ -143,6 +144,17 @@
       }
       break;
     }
+    case webrtc::kVideoCodecAV1: {
+      absl::optional<webrtc::AV1Profile> profile =
+          webrtc::ParseSdpForAV1Profile(format.parameters);
+      bool lossless_color = profile.has_value() &&
+                            profile.value() == webrtc::AV1Profile::kProfile1;
+      VLOG(0) << "Creating AV1 encoder, lossless_color="
+              << (lossless_color ? "true" : "false");
+      encoder_ = std::make_unique<WebrtcVideoEncoderAV1>();
+      encoder_->SetLosslessColor(lossless_color);
+      break;
+    }
     case webrtc::kVideoCodecH264:
 #if defined(USE_H264_ENCODER)
       VLOG(0) << "Creating H264 encoder.";
@@ -151,10 +163,6 @@
       NOTIMPLEMENTED();
 #endif
       break;
-    case webrtc::kVideoCodecAV1:
-      VLOG(0) << "Creating AV1 encoder.";
-      encoder_ = std::make_unique<WebrtcVideoEncoderAV1>();
-      break;
     default:
       LOG(FATAL) << "Unknown codec type: " << codec_type_;
   }