yolobs-studio/plugins/mac-capture/mac-window-capture.m

#include <obs-module.h>
#include <util/darray.h>
#include <util/threading.h>
#include <util/platform.h>

#import <CoreGraphics/CGWindow.h>
#import <Cocoa/Cocoa.h>

#include "window-utils.h"

struct window_capture {
	obs_source_t *source;

	struct cocoa_window window;

	//CGRect              bounds;
	//CGWindowListOption  window_option;
	CGWindowImageOption image_option;

	CGColorSpaceRef color_space;

	DARRAY(uint8_t) buffer;

	pthread_t  capture_thread;
	os_event_t *capture_event;
	os_event_t *stop_event;
};

static CGImageRef get_image(struct window_capture *wc)
{
	NSArray *arr = (NSArray*)CGWindowListCreate(
			kCGWindowListOptionIncludingWindow,
			wc->window.window_id);
	[arr autorelease];

	if (arr.count)
		return CGWindowListCreateImage(CGRectNull,
				kCGWindowListOptionIncludingWindow,
				wc->window.window_id, wc->image_option);

	if (!find_window(&wc->window, NULL, false))
		return NULL;

	return CGWindowListCreateImage(CGRectNull,
			kCGWindowListOptionIncludingWindow,
			wc->window.window_id, wc->image_option);
}

static inline void capture_frame(struct window_capture *wc)
{
	uint64_t   ts  = os_gettime_ns();
	CGImageRef img = get_image(wc);
	if (!img)
		return;

	size_t width  = CGImageGetWidth(img);
	size_t height = CGImageGetHeight(img);

	CGRect rect = {{0, 0}, {width, height}};
	da_reserve(wc->buffer, width * height * 4);
	uint8_t *data = wc->buffer.array;

	CGContextRef cg_context = CGBitmapContextCreate(data, width, height,
			8, width * 4, wc->color_space,
			kCGBitmapByteOrder32Host |
			kCGImageAlphaPremultipliedFirst);
	CGContextSetBlendMode(cg_context, kCGBlendModeCopy);
	CGContextDrawImage(cg_context, rect, img);
	CGContextRelease(cg_context);
	CGImageRelease(img);

	struct obs_source_frame frame = {
		.format      = VIDEO_FORMAT_BGRA,
		.width       = width,
		.height      = height,
		.data[0]     = data,
		.linesize[0] = width * 4,
		.timestamp   = ts,
	};

	obs_source_output_video(wc->source, &frame);
}

static void *capture_thread(void *data)
{
	struct window_capture *wc = data;

	for (;;) {
		os_event_wait(wc->capture_event);
		if (os_event_try(wc->stop_event) != EAGAIN)
			break;

		@autoreleasepool {
			capture_frame(wc);
		}
	}

	return NULL;
}

static inline void *window_capture_create_internal(obs_data_t *settings,
		obs_source_t *source)
{
	struct window_capture *wc = bzalloc(sizeof(struct window_capture));

	wc->source = source;

	wc->color_space = CGColorSpaceCreateDeviceRGB();

	da_init(wc->buffer);

	init_window(&wc->window, settings);

	wc->image_option = obs_data_get_bool(settings, "show_shadow") ?
		kCGWindowImageDefault : kCGWindowImageBoundsIgnoreFraming;

	os_event_init(&wc->capture_event, OS_EVENT_TYPE_AUTO);
	os_event_init(&wc->stop_event, OS_EVENT_TYPE_MANUAL);

	pthread_create(&wc->capture_thread, NULL, capture_thread, wc);

	return wc;
}

static void *window_capture_create(obs_data_t *settings, obs_source_t *source)
{
	@autoreleasepool {
		return window_capture_create_internal(settings, source);
	}
}

static void window_capture_destroy(void *data)
{
	struct window_capture *cap = data;

	os_event_signal(cap->stop_event);
	os_event_signal(cap->capture_event);
	
	pthread_join(cap->capture_thread, NULL);

	CGColorSpaceRelease(cap->color_space);

	da_free(cap->buffer);

	os_event_destroy(cap->capture_event);
	os_event_destroy(cap->stop_event);

	destroy_window(&cap->window);

	bfree(cap);
}

static void window_capture_defaults(obs_data_t *settings)
{
	obs_data_set_default_bool(settings, "show_shadow", false);
	window_defaults(settings);
}

static obs_properties_t *window_capture_properties(void *unused)
{
	UNUSED_PARAMETER(unused);

	obs_properties_t *props = obs_properties_create();

	add_window_properties(props);

	obs_properties_add_bool(props, "show_shadow",
			obs_module_text("WindowCapture.ShowShadow"));

	return props;
}

static inline void window_capture_update_internal(struct window_capture *wc,
		obs_data_t *settings)
{
	wc->image_option = obs_data_get_bool(settings, "show_shadow") ?
		kCGWindowImageDefault : kCGWindowImageBoundsIgnoreFraming;

	update_window(&wc->window, settings);
}

static void window_capture_update(void *data, obs_data_t *settings)
{
	@autoreleasepool {
		return window_capture_update_internal(data, settings);
	}
}

static const char *window_capture_getname(void *unused)
{
	UNUSED_PARAMETER(unused);
	return obs_module_text("WindowCapture");
}

static inline void window_capture_tick_internal(struct window_capture *wc,
		float seconds)
{
	UNUSED_PARAMETER(seconds);
	os_event_signal(wc->capture_event);
}

static void window_capture_tick(void *data, float seconds)
{
	struct window_capture *wc = data;

	if (!obs_source_showing(wc->source))
		return;

	@autoreleasepool {
		return window_capture_tick_internal(data, seconds);
	}
}

struct obs_source_info window_capture_info = {
	.id             = "window_capture",
	.type           = OBS_SOURCE_TYPE_INPUT,
	.get_name       = window_capture_getname,

	.create         = window_capture_create,
	.destroy        = window_capture_destroy,

	.output_flags   = OBS_SOURCE_ASYNC_VIDEO,
	.video_tick     = window_capture_tick,

	.get_defaults   = window_capture_defaults,
	.get_properties = window_capture_properties,
	.update         = window_capture_update,
};