diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCapture.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCapture.kt
new file mode 100644
index 0000000..39f35a5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCapture.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * 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.
+ */
+package com.android.systemui.screenshot
+
+import android.graphics.Bitmap
+import android.graphics.Rect
+
+interface ImageCapture {
+
+    fun captureDisplay(displayId: Int, crop: Rect? = null): Bitmap?
+
+    fun captureTask(taskId: Int): Bitmap?
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
new file mode 100644
index 0000000..258c436
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
@@ -0,0 +1,78 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * 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.
+ */
+
+package com.android.systemui.screenshot
+
+import android.app.IActivityTaskManager
+import android.graphics.Bitmap
+import android.graphics.Rect
+import android.hardware.display.DisplayManager
+import android.os.IBinder
+import android.util.Log
+import android.view.DisplayAddress
+import android.view.SurfaceControl
+import android.view.SurfaceControl.DisplayCaptureArgs
+import android.view.SurfaceControl.ScreenshotHardwareBuffer
+import androidx.annotation.VisibleForTesting
+import javax.inject.Inject
+
+private const val TAG = "ImageCaptureImpl"
+
+open class ImageCaptureImpl @Inject constructor(
+    private val displayManager: DisplayManager,
+    private val atmService: IActivityTaskManager
+) : ImageCapture {
+
+    override fun captureDisplay(displayId: Int, crop: Rect?): Bitmap? {
+        val width = crop?.width() ?: 0
+        val height = crop?.height() ?: 0
+        val sourceCrop = crop ?: Rect()
+        val displayToken = physicalDisplayToken(displayId) ?: return null
+        val buffer = captureDisplay(displayToken, width, height, sourceCrop)
+
+        return buffer?.asBitmap()
+    }
+
+    override fun captureTask(taskId: Int): Bitmap? {
+        val snapshot = atmService.takeTaskSnapshot(taskId)
+        return Bitmap.wrapHardwareBuffer(snapshot.hardwareBuffer, snapshot.colorSpace)
+    }
+
+    @VisibleForTesting
+    open fun physicalDisplayToken(displayId: Int): IBinder? {
+        val display = displayManager.getDisplay(displayId)
+        if (display == null) {
+            Log.e(TAG, "No display with id: $displayId")
+            return null
+        }
+        val address = display.address
+        if (address !is DisplayAddress.Physical) {
+            Log.e(TAG, "Display does not have a physical address: $display")
+            return null
+        }
+        return SurfaceControl.getPhysicalDisplayToken(address.physicalDisplayId)
+    }
+
+    @VisibleForTesting
+    open fun captureDisplay(displayToken: IBinder, width: Int, height: Int, crop: Rect): ScreenshotHardwareBuffer? {
+        val captureArgs = DisplayCaptureArgs.Builder(displayToken)
+            .setSize(width, height)
+            .setSourceCrop(crop)
+            .build()
+        return SurfaceControl.captureDisplay(captureArgs)
+    }
+
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index 82de389..69ee8e8 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -57,14 +57,12 @@
 import android.media.MediaPlayer;
 import android.net.Uri;
 import android.os.Bundle;
-import android.os.IBinder;
 import android.os.RemoteException;
 import android.provider.Settings;
 import android.util.DisplayMetrics;
 import android.util.Log;
 import android.util.Pair;
 import android.view.Display;
-import android.view.DisplayAddress;
 import android.view.IRemoteAnimationFinishedCallback;
 import android.view.IRemoteAnimationRunner;
 import android.view.KeyEvent;
@@ -72,7 +70,6 @@
 import android.view.RemoteAnimationAdapter;
 import android.view.RemoteAnimationTarget;
 import android.view.ScrollCaptureResponse;
-import android.view.SurfaceControl;
 import android.view.View;
 import android.view.ViewRootImpl;
 import android.view.ViewTreeObserver;
@@ -249,6 +246,7 @@
     private final ScreenshotSmartActions mScreenshotSmartActions;
     private final UiEventLogger mUiEventLogger;
     private final ImageExporter mImageExporter;
+    private final ImageCapture mImageCapture;
     private final Executor mMainExecutor;
     private final ExecutorService mBgExecutor;
     private final BroadcastSender mBroadcastSender;
@@ -295,6 +293,7 @@
             ScrollCaptureClient scrollCaptureClient,
             UiEventLogger uiEventLogger,
             ImageExporter imageExporter,
+            ImageCapture imageCapture,
             @Main Executor mainExecutor,
             ScrollCaptureController scrollCaptureController,
             LongScreenshotData longScreenshotHolder,
@@ -308,6 +307,7 @@
         mScrollCaptureClient = scrollCaptureClient;
         mUiEventLogger = uiEventLogger;
         mImageExporter = imageExporter;
+        mImageCapture = imageCapture;
         mMainExecutor = mainExecutor;
         mScrollCaptureController = scrollCaptureController;
         mLongScreenshotHolder = longScreenshotHolder;
@@ -531,7 +531,7 @@
 
         // copy the input Rect, since SurfaceControl.screenshot can mutate it
         Rect screenRect = new Rect(crop);
-        Bitmap screenshot = captureScreenshot(crop);
+        Bitmap screenshot = mImageCapture.captureDisplay(DEFAULT_DISPLAY, crop);
 
         if (screenshot == null) {
             Log.e(TAG, "takeScreenshotInternal: Screenshot bitmap was null");
@@ -549,32 +549,6 @@
                 ClipboardOverlayController.SELF_PERMISSION);
     }
 
-    private Bitmap captureScreenshot(Rect crop) {
-        int width = crop.width();
-        int height = crop.height();
-        Bitmap screenshot = null;
-        final Display display = getDefaultDisplay();
-        final DisplayAddress address = display.getAddress();
-        if (!(address instanceof DisplayAddress.Physical)) {
-            Log.e(TAG, "Skipping Screenshot - Default display does not have a physical address: "
-                    + display);
-        } else {
-            final DisplayAddress.Physical physicalAddress = (DisplayAddress.Physical) address;
-
-            final IBinder displayToken = SurfaceControl.getPhysicalDisplayToken(
-                    physicalAddress.getPhysicalDisplayId());
-            final SurfaceControl.DisplayCaptureArgs captureArgs =
-                    new SurfaceControl.DisplayCaptureArgs.Builder(displayToken)
-                            .setSourceCrop(crop)
-                            .setSize(width, height)
-                            .build();
-            final SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
-                    SurfaceControl.captureDisplay(captureArgs);
-            screenshot = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
-        }
-        return screenshot;
-    }
-
     private void saveScreenshot(Bitmap screenshot, Consumer<Uri> finisher, Rect screenRect,
             Insets screenInsets, ComponentName topComponent, boolean showFlash) {
         withWindowAttached(() ->
@@ -720,7 +694,7 @@
             mScreenshotView.showScrollChip(response.getPackageName(), /* onClick */ () -> {
                 DisplayMetrics displayMetrics = new DisplayMetrics();
                 getDefaultDisplay().getRealMetrics(displayMetrics);
-                Bitmap newScreenshot = captureScreenshot(
+                Bitmap newScreenshot = mImageCapture.captureDisplay(DEFAULT_DISPLAY,
                         new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels));
 
                 mScreenshotView.prepareScrollingTransition(response, mScreenBitmap, newScreenshot,
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
index 91ef3c3..3e442587 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
@@ -18,6 +18,8 @@
 
 import android.app.Service;
 
+import com.android.systemui.screenshot.ImageCapture;
+import com.android.systemui.screenshot.ImageCaptureImpl;
 import com.android.systemui.screenshot.TakeScreenshotService;
 
 import dagger.Binds;
@@ -37,4 +39,6 @@
     @ClassKey(TakeScreenshotService.class)
     public abstract Service bindTakeScreenshotService(TakeScreenshotService service);
 
+    @Binds
+    public abstract ImageCapture bindImageCapture(ImageCaptureImpl capture);
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt
new file mode 100644
index 0000000..ce3f20d
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt
@@ -0,0 +1,84 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * 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.
+ */
+
+package com.android.systemui.screenshot
+
+import android.app.IActivityTaskManager
+import android.graphics.Rect
+import android.hardware.display.DisplayManager
+import android.os.Binder
+import android.os.IBinder
+import android.testing.AndroidTestingRunner
+import android.view.Display
+import android.view.SurfaceControl.ScreenshotHardwareBuffer
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.mock
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Test the logic within ImageCaptureImpl
+ */
+@RunWith(AndroidTestingRunner::class)
+class ImageCaptureImplTest : SysuiTestCase() {
+    private val displayManager = mock<DisplayManager>()
+    private val atmService = mock<IActivityTaskManager>()
+    private val capture = TestableImageCaptureImpl(displayManager, atmService)
+
+    @Test
+    fun captureDisplayWithCrop() {
+        capture.captureDisplay(Display.DEFAULT_DISPLAY, Rect(1, 2, 3, 4))
+        assertThat(capture.token).isNotNull()
+        assertThat(capture.width!!).isEqualTo(2)
+        assertThat(capture.height!!).isEqualTo(2)
+        assertThat(capture.crop!!).isEqualTo(Rect(1, 2, 3, 4))
+    }
+
+    @Test
+    fun captureDisplayWithNullCrop() {
+        capture.captureDisplay(Display.DEFAULT_DISPLAY, null)
+        assertThat(capture.token).isNotNull()
+        assertThat(capture.width!!).isEqualTo(0)
+        assertThat(capture.height!!).isEqualTo(0)
+        assertThat(capture.crop!!).isEqualTo(Rect())
+    }
+
+    class TestableImageCaptureImpl(
+        displayManager: DisplayManager,
+        atmService: IActivityTaskManager
+    ) :
+        ImageCaptureImpl(displayManager, atmService) {
+
+        var token: IBinder? = null
+        var width: Int? = null
+        var height: Int? = null
+        var crop: Rect? = null
+
+        override fun physicalDisplayToken(displayId: Int): IBinder {
+            return Binder()
+        }
+
+        override fun captureDisplay(displayToken: IBinder, width: Int, height: Int, crop: Rect):
+                ScreenshotHardwareBuffer {
+            this.token = displayToken
+            this.width = width
+            this.height = height
+            this.crop = crop
+            return ScreenshotHardwareBuffer(null, null, false, false)
+        }
+    }
+}
\ No newline at end of file