diff --git a/core/java/com/android/internal/util/ScreenshotHelper.java b/core/java/com/android/internal/util/ScreenshotHelper.java
index 8af2450f29efa..b9243ece91585 100644
--- a/core/java/com/android/internal/util/ScreenshotHelper.java
+++ b/core/java/com/android/internal/util/ScreenshotHelper.java
@@ -219,8 +219,9 @@ public class ScreenshotHelper {
                 throw new IllegalArgumentException("Bundle does not contain a hardware bitmap");
             }
 
-            HardwareBuffer buffer = bundle.getParcelable(KEY_BUFFER);
-            ParcelableColorSpace colorSpace = bundle.getParcelable(KEY_COLOR_SPACE);
+            HardwareBuffer buffer = bundle.getParcelable(KEY_BUFFER, HardwareBuffer.class);
+            ParcelableColorSpace colorSpace = bundle.getParcelable(KEY_COLOR_SPACE,
+                    ParcelableColorSpace.class);
 
             return Bitmap.wrapHardwareBuffer(Objects.requireNonNull(buffer),
                     colorSpace.getColorSpace());