Merge pull request #75142 from bruvzg/scr

[DisplayServer] Implement screen_get_image method for LinuxBSD/X11, macOS and Windows.

9 files changed, 232 insertions, 4 deletions

diff --git a/doc/classes/DisplayServer.xml b/doc/classes/DisplayServer.xml
index 513414e695..beccaab048 100644
--- a/doc/classes/DisplayServer.xml
+++ b/doc/classes/DisplayServer.xml
@@ -855,6 +855,15 @@
 				[b]Note:[/b] This method is implemented on Android, Linux (X11), macOS and Windows. Returns [code]72[/code] on unsupported platforms.
 			</description>
 		</method>
+		<method name="screen_get_image" qualifiers="const">
+			<return type="Image" />
+			<param index="0" name="screen" type="int" default="-1" />
+			<description>
+				Returns screenshot of the [param screen].
+				[b]Note:[/b] This method is implemented on Linux (X11), macOS, and Windows.
+				[b]Note:[/b] On macOS, this method requires "Screen Recording" permission, if permission is not granted it will return desktop wallpaper color.
+			</description>
+		</method>
 		<method name="screen_get_max_scale" qualifiers="const">
 			<return type="float" />
 			<description>
diff --git a/platform/linuxbsd/x11/display_server_x11.cpp b/platform/linuxbsd/x11/display_server_x11.cpp
index 13f261aa43..9506189573 100644
--- a/platform/linuxbsd/x11/display_server_x11.cpp
+++ b/platform/linuxbsd/x11/display_server_x11.cpp
@@ -1206,6 +1206,105 @@ Color DisplayServerX11::screen_get_pixel(const Point2i &p_position) const {
 	return Color();
 }
 
+Ref<Image> DisplayServerX11::screen_get_image(int p_screen) const {
+	ERR_FAIL_INDEX_V(p_screen, get_screen_count(), Ref<Image>());
+
+	switch (p_screen) {
+		case SCREEN_PRIMARY: {
+			p_screen = get_primary_screen();
+		} break;
+		case SCREEN_OF_MAIN_WINDOW: {
+			p_screen = window_get_current_screen(MAIN_WINDOW_ID);
+		} break;
+		default:
+			break;
+	}
+
+	ERR_FAIL_COND_V(p_screen < 0, Ref<Image>());
+
+	XImage *image = nullptr;
+
+	int event_base, error_base;
+	if (XineramaQueryExtension(x11_display, &event_base, &error_base)) {
+		int xin_count;
+		XineramaScreenInfo *xsi = XineramaQueryScreens(x11_display, &xin_count);
+		if (p_screen < xin_count) {
+			int x_count = XScreenCount(x11_display);
+			for (int i = 0; i < x_count; i++) {
+				Window root = XRootWindow(x11_display, i);
+				XWindowAttributes root_attrs;
+				XGetWindowAttributes(x11_display, root, &root_attrs);
+				if ((xsi[p_screen].x_org >= root_attrs.x) && (xsi[p_screen].x_org <= root_attrs.x + root_attrs.width) && (xsi[p_screen].y_org >= root_attrs.y) && (xsi[p_screen].y_org <= root_attrs.y + root_attrs.height)) {
+					image = XGetImage(x11_display, root, xsi[p_screen].x_org, xsi[p_screen].y_org, xsi[p_screen].width, xsi[p_screen].height, AllPlanes, ZPixmap);
+					break;
+				}
+			}
+		} else {
+			ERR_FAIL_V_MSG(Ref<Image>(), "Invalid screen index: " + itos(p_screen) + "(count: " + itos(xin_count) + ").");
+		}
+	} else {
+		int x_count = XScreenCount(x11_display);
+		if (p_screen < x_count) {
+			Window root = XRootWindow(x11_display, p_screen);
+
+			XWindowAttributes root_attrs;
+			XGetWindowAttributes(x11_display, root, &root_attrs);
+
+			image = XGetImage(x11_display, root, root_attrs.x, root_attrs.y, root_attrs.width, root_attrs.height, AllPlanes, ZPixmap);
+		} else {
+			ERR_FAIL_V_MSG(Ref<Image>(), "Invalid screen index: " + itos(p_screen) + "(count: " + itos(x_count) + ").");
+		}
+	}
+
+	Ref<Image> img;
+	if (image) {
+		int width = image->width;
+		int height = image->height;
+
+		Vector<uint8_t> img_data;
+		img_data.resize(height * width * 4);
+
+		uint8_t *sr = (uint8_t *)image->data;
+		uint8_t *wr = (uint8_t *)img_data.ptrw();
+
+		if (image->bits_per_pixel == 24 && image->red_mask == 0xff0000 && image->green_mask == 0x00ff00 && image->blue_mask == 0x0000ff) {
+			for (int y = 0; y < height; y++) {
+				for (int x = 0; x < width; x++) {
+					wr[(y * width + x) * 4 + 0] = sr[(y * width + x) * 3 + 2];
+					wr[(y * width + x) * 4 + 1] = sr[(y * width + x) * 3 + 1];
+					wr[(y * width + x) * 4 + 2] = sr[(y * width + x) * 3 + 0];
+					wr[(y * width + x) * 4 + 3] = 255;
+				}
+			}
+		} else if (image->bits_per_pixel == 24 && image->red_mask == 0x0000ff && image->green_mask == 0x00ff00 && image->blue_mask == 0xff0000) {
+			for (int y = 0; y < height; y++) {
+				for (int x = 0; x < width; x++) {
+					wr[(y * width + x) * 4 + 0] = sr[(y * width + x) * 3 + 2];
+					wr[(y * width + x) * 4 + 1] = sr[(y * width + x) * 3 + 1];
+					wr[(y * width + x) * 4 + 2] = sr[(y * width + x) * 3 + 0];
+					wr[(y * width + x) * 4 + 3] = 255;
+				}
+			}
+		} else if (image->bits_per_pixel == 32 && image->red_mask == 0xff0000 && image->green_mask == 0x00ff00 && image->blue_mask == 0x0000ff) {
+			for (int y = 0; y < height; y++) {
+				for (int x = 0; x < width; x++) {
+					wr[(y * width + x) * 4 + 0] = sr[(y * width + x) * 4 + 2];
+					wr[(y * width + x) * 4 + 1] = sr[(y * width + x) * 4 + 1];
+					wr[(y * width + x) * 4 + 2] = sr[(y * width + x) * 4 + 0];
+					wr[(y * width + x) * 4 + 3] = 255;
+				}
+			}
+		} else {
+			XFree(image);
+			ERR_FAIL_V_MSG(Ref<Image>(), vformat("XImage with RGB mask %x %x %x and depth %d is not supported.", image->red_mask, image->green_mask, image->blue_mask, image->bits_per_pixel));
+		}
+		img = Image::create_from_data(width, height, false, Image::FORMAT_RGBA8, img_data);
+		XFree(image);
+	}
+
+	return img;
+}
+
 float DisplayServerX11::screen_get_refresh_rate(int p_screen) const {
 	_THREAD_SAFE_METHOD_
 
diff --git a/platform/linuxbsd/x11/display_server_x11.h b/platform/linuxbsd/x11/display_server_x11.h
index 6d343be3ab..fd3a5dccfa 100644
--- a/platform/linuxbsd/x11/display_server_x11.h
+++ b/platform/linuxbsd/x11/display_server_x11.h
@@ -412,6 +412,7 @@ public:
 	virtual int screen_get_dpi(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual float screen_get_refresh_rate(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual Color screen_get_pixel(const Point2i &p_position) const override;
+	virtual Ref<Image> screen_get_image(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 
 #if defined(DBUS_ENABLED)
 	virtual void screen_set_keep_on(bool p_enable) override;
diff --git a/platform/macos/display_server_macos.h b/platform/macos/display_server_macos.h
index a1cd83280d..495dc43c55 100644
--- a/platform/macos/display_server_macos.h
+++ b/platform/macos/display_server_macos.h
@@ -336,6 +336,7 @@ public:
 	virtual Rect2i screen_get_usable_rect(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual float screen_get_refresh_rate(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual Color screen_get_pixel(const Point2i &p_position) const override;
+	virtual Ref<Image> screen_get_image(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual void screen_set_keep_on(bool p_enable) override;
 	virtual bool screen_is_kept_on() const override;
 
diff --git a/platform/macos/display_server_macos.mm b/platform/macos/display_server_macos.mm
index 32f9441484..4657461dca 100644
--- a/platform/macos/display_server_macos.mm
+++ b/platform/macos/display_server_macos.mm
@@ -2234,6 +2234,49 @@ Color DisplayServerMacOS::screen_get_pixel(const Point2i &p_position) const {
 	return Color();
 }
 
+Ref<Image> DisplayServerMacOS::screen_get_image(int p_screen) const {
+	ERR_FAIL_INDEX_V(p_screen, get_screen_count(), Ref<Image>());
+
+	switch (p_screen) {
+		case SCREEN_PRIMARY: {
+			p_screen = get_primary_screen();
+		} break;
+		case SCREEN_OF_MAIN_WINDOW: {
+			p_screen = window_get_current_screen(MAIN_WINDOW_ID);
+		} break;
+		default:
+			break;
+	}
+
+	Ref<Image> img;
+	NSArray *screenArray = [NSScreen screens];
+	if ((NSUInteger)p_screen < [screenArray count]) {
+		NSRect nsrect = [[screenArray objectAtIndex:p_screen] frame];
+		NSDictionary *screenDescription = [[screenArray objectAtIndex:p_screen] deviceDescription];
+		CGDirectDisplayID display_id = [[screenDescription objectForKey:@"NSScreenNumber"] unsignedIntValue];
+		CGImageRef image = CGDisplayCreateImageForRect(display_id, CGRectMake(0, 0, nsrect.size.width, nsrect.size.height));
+		if (image) {
+			CGColorSpaceRef color_space = CGColorSpaceCreateDeviceRGB();
+			if (color_space) {
+				NSUInteger width = CGImageGetWidth(image);
+				NSUInteger height = CGImageGetHeight(image);
+
+				Vector<uint8_t> img_data;
+				img_data.resize(height * width * 4);
+				CGContextRef context = CGBitmapContextCreate(img_data.ptrw(), width, height, 8, 4 * width, color_space, kCGImageAlphaPremultipliedLast | kCGBitmapByteOrder32Big);
+				if (context) {
+					CGContextDrawImage(context, CGRectMake(0, 0, width, height), image);
+					img = Image::create_from_data(width, height, false, Image::FORMAT_RGBA8, img_data);
+					CGContextRelease(context);
+				}
+				CGColorSpaceRelease(color_space);
+			}
+			CGImageRelease(image);
+		}
+	}
+	return img;
+}
+
 float DisplayServerMacOS::screen_get_refresh_rate(int p_screen) const {
 	_THREAD_SAFE_METHOD_
 
diff --git a/platform/windows/display_server_windows.cpp b/platform/windows/display_server_windows.cpp
index 02bb30ed38..ff32a318a2 100644
--- a/platform/windows/display_server_windows.cpp
+++ b/platform/windows/display_server_windows.cpp
@@ -614,15 +614,87 @@ Color DisplayServerWindows::screen_get_pixel(const Point2i &p_position) const {
 		win81p_LogicalToPhysicalPointForPerMonitorDPI(0, &p);
 	}
 	HDC dc = GetDC(0);
-	COLORREF col = GetPixel(dc, p.x, p.y);
-	if (col != CLR_INVALID) {
-		return Color(float(col & 0x000000FF) / 256.0, float((col & 0x0000FF00) >> 8) / 256.0, float((col & 0x00FF0000) >> 16) / 256.0, 1.0);
+	if (dc) {
+		COLORREF col = GetPixel(dc, p.x, p.y);
+		if (col != CLR_INVALID) {
+			ReleaseDC(NULL, dc);
+			return Color(float(col & 0x000000FF) / 256.0, float((col & 0x0000FF00) >> 8) / 256.0, float((col & 0x00FF0000) >> 16) / 256.0, 1.0);
+		}
+		ReleaseDC(NULL, dc);
 	}
-	ReleaseDC(NULL, dc);
 
 	return Color();
 }
 
+Ref<Image> DisplayServerWindows::screen_get_image(int p_screen) const {
+	ERR_FAIL_INDEX_V(p_screen, get_screen_count(), Ref<Image>());
+
+	switch (p_screen) {
+		case SCREEN_PRIMARY: {
+			p_screen = get_primary_screen();
+		} break;
+		case SCREEN_OF_MAIN_WINDOW: {
+			p_screen = window_get_current_screen(MAIN_WINDOW_ID);
+		} break;
+		default:
+			break;
+	}
+
+	Point2i pos = screen_get_position(p_screen) + _get_screens_origin();
+	Size2i size = screen_get_size(p_screen);
+
+	POINT p1;
+	p1.x = pos.x;
+	p1.y = pos.y;
+
+	POINT p2;
+	p2.x = pos.x + size.x;
+	p2.y = pos.y + size.y;
+	if (win81p_LogicalToPhysicalPointForPerMonitorDPI) {
+		win81p_LogicalToPhysicalPointForPerMonitorDPI(0, &p1);
+		win81p_LogicalToPhysicalPointForPerMonitorDPI(0, &p2);
+	}
+
+	Ref<Image> img;
+	HDC dc = GetDC(0);
+	if (dc) {
+		HDC hdc = CreateCompatibleDC(dc);
+		int width = p2.x - p1.x;
+		int height = p2.y - p1.y;
+		if (hdc) {
+			HBITMAP hbm = CreateCompatibleBitmap(dc, width, height);
+			if (hbm) {
+				SelectObject(hdc, hbm);
+				BitBlt(hdc, 0, 0, width, height, dc, p1.x, p1.y, SRCCOPY);
+
+				BITMAPINFO bmp_info = { 0 };
+				bmp_info.bmiHeader.biSize = sizeof(bmp_info.bmiHeader);
+				bmp_info.bmiHeader.biWidth = width;
+				bmp_info.bmiHeader.biHeight = -height;
+				bmp_info.bmiHeader.biPlanes = 1;
+				bmp_info.bmiHeader.biBitCount = 32;
+				bmp_info.bmiHeader.biCompression = BI_RGB;
+
+				Vector<uint8_t> img_data;
+				img_data.resize(width * height * 4);
+				GetDIBits(hdc, hbm, 0, height, img_data.ptrw(), &bmp_info, DIB_RGB_COLORS);
+
+				uint8_t *wr = (uint8_t *)img_data.ptrw();
+				for (int i = 0; i < width * height; i++) {
+					SWAP(wr[i * 4 + 0], wr[i * 4 + 2]);
+				}
+				img = Image::create_from_data(width, height, false, Image::FORMAT_RGBA8, img_data);
+
+				DeleteObject(hbm);
+			}
+			DeleteDC(hdc);
+		}
+		ReleaseDC(NULL, dc);
+	}
+
+	return img;
+}
+
 float DisplayServerWindows::screen_get_refresh_rate(int p_screen) const {
 	_THREAD_SAFE_METHOD_
 
diff --git a/platform/windows/display_server_windows.h b/platform/windows/display_server_windows.h
index a8dcc6bad0..4a1619f331 100644
--- a/platform/windows/display_server_windows.h
+++ b/platform/windows/display_server_windows.h
@@ -529,6 +529,7 @@ public:
 	virtual int screen_get_dpi(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual float screen_get_refresh_rate(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 	virtual Color screen_get_pixel(const Point2i &p_position) const override;
+	virtual Ref<Image> screen_get_image(int p_screen = SCREEN_OF_MAIN_WINDOW) const override;
 
 	virtual void screen_set_keep_on(bool p_enable) override; //disable screensaver
 	virtual bool screen_is_kept_on() const override;
diff --git a/servers/display_server.cpp b/servers/display_server.cpp
index dc8c308389..516f2f1be6 100644
--- a/servers/display_server.cpp
+++ b/servers/display_server.cpp
@@ -660,6 +660,7 @@ void DisplayServer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("screen_get_max_scale"), &DisplayServer::screen_get_max_scale);
 	ClassDB::bind_method(D_METHOD("screen_get_refresh_rate", "screen"), &DisplayServer::screen_get_refresh_rate, DEFVAL(SCREEN_OF_MAIN_WINDOW));
 	ClassDB::bind_method(D_METHOD("screen_get_pixel", "position"), &DisplayServer::screen_get_pixel);
+	ClassDB::bind_method(D_METHOD("screen_get_image", "screen"), &DisplayServer::screen_get_image, DEFVAL(SCREEN_OF_MAIN_WINDOW));
 
 	ClassDB::bind_method(D_METHOD("screen_set_orientation", "orientation", "screen"), &DisplayServer::screen_set_orientation, DEFVAL(SCREEN_OF_MAIN_WINDOW));
 	ClassDB::bind_method(D_METHOD("screen_get_orientation", "screen"), &DisplayServer::screen_get_orientation, DEFVAL(SCREEN_OF_MAIN_WINDOW));
diff --git a/servers/display_server.h b/servers/display_server.h
index 8f2df62094..d8e67b4f92 100644
--- a/servers/display_server.h
+++ b/servers/display_server.h
@@ -301,6 +301,7 @@ public:
 	}
 	virtual float screen_get_refresh_rate(int p_screen = SCREEN_OF_MAIN_WINDOW) const = 0;
 	virtual Color screen_get_pixel(const Point2i &p_position) const { return Color(); };
+	virtual Ref<Image> screen_get_image(int p_screen = SCREEN_OF_MAIN_WINDOW) const { return Ref<Image>(); };
 	virtual bool is_touchscreen_available() const;
 
 	// Keep the ScreenOrientation enum values in sync with the `display/window/handheld/orientation`