feat(rembg): add RemoveBG Studio - native background removal app
C++ WebView2 host app driving a Python/rembg engine as a JSON-stdio sidecar, with live before/after preview for images and video frames. Also includes the original CLI (removebg.py) built on the same core/ engine. The WebView2 SDK is fetched on demand via native/third_party/fetch-webview2.ps1 (wired into the Makefile) rather than vendored, since it's ~15MB of prebuilt/generated Microsoft SDK content. A Makefile provides make sync/build/run/clean as the standard entry points.
This commit is contained in:
@@ -0,0 +1,237 @@
|
||||
#include <windows.h>
|
||||
#include <dwmapi.h>
|
||||
#include <shellapi.h>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <nlohmann/json.hpp>
|
||||
|
||||
#include "WebViewHost.h"
|
||||
#include "SidecarProcess.h"
|
||||
#include "Dialogs.h"
|
||||
#include "PathUtil.h"
|
||||
|
||||
#pragma comment(lib, "dwmapi.lib")
|
||||
|
||||
using json = nlohmann::json;
|
||||
using namespace rbg;
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr wchar_t kWindowClass[] = L"RemoveBGStudioWindow";
|
||||
constexpr wchar_t kWindowTitle[] = L"RemoveBG Studio";
|
||||
constexpr UINT WM_APP_SIDECAR_EVENT = WM_APP + 1;
|
||||
|
||||
#ifndef DWMWA_USE_IMMERSIVE_DARK_MODE
|
||||
#define DWMWA_USE_IMMERSIVE_DARK_MODE 20
|
||||
#endif
|
||||
|
||||
struct AppState {
|
||||
HWND hwnd = nullptr;
|
||||
std::unique_ptr<WebViewHost> webView;
|
||||
std::unique_ptr<SidecarProcess> sidecar;
|
||||
};
|
||||
|
||||
AppState* g_app = nullptr;
|
||||
|
||||
std::vector<FilterSpec> ImageFilters() {
|
||||
return { { L"Image Files", L"*.png;*.jpg;*.jpeg;*.webp;*.bmp;*.tiff" }, { L"All Files", L"*.*" } };
|
||||
}
|
||||
|
||||
std::vector<FilterSpec> VideoFilters() {
|
||||
return { { L"Video Files", L"*.mp4;*.mov;*.avi;*.mkv;*.webm;*.flv" }, { L"All Files", L"*.*" } };
|
||||
}
|
||||
|
||||
// Handles JS messages addressed to the native host directly (file/folder/
|
||||
// color dialogs) rather than the Python sidecar. Replies synchronously via
|
||||
// PostWebMessageAsJson since we're already on the UI thread here.
|
||||
void HandleNativeCommand(AppState& app, const json& req) {
|
||||
int id = req.value("id", 0);
|
||||
std::string cmd = req.value("cmd", "");
|
||||
json params = req.value("params", json::object());
|
||||
|
||||
json reply = { {"id", id} };
|
||||
|
||||
auto ok = [&](json result) {
|
||||
reply["event"] = "done";
|
||||
reply["result"] = std::move(result);
|
||||
};
|
||||
auto fail = [&](const std::string& msg) {
|
||||
reply["event"] = "error";
|
||||
reply["message"] = msg;
|
||||
};
|
||||
|
||||
if (cmd == "dialog-open-image") {
|
||||
auto path = ShowOpenFileDialog(app.hwnd, L"Select Image File", ImageFilters());
|
||||
ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
|
||||
} else if (cmd == "dialog-open-images") {
|
||||
auto paths = ShowOpenMultiFileDialog(app.hwnd, L"Select Images", ImageFilters());
|
||||
json arr = json::array();
|
||||
for (auto& p : paths) arr.push_back(WideToUtf8(p));
|
||||
ok(json{ {"paths", arr} });
|
||||
} else if (cmd == "dialog-open-folder") {
|
||||
std::wstring title = Utf8ToWide(params.value("title", std::string("Select Folder")));
|
||||
auto path = ShowOpenFolderDialog(app.hwnd, title);
|
||||
ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
|
||||
} else if (cmd == "dialog-open-video") {
|
||||
auto path = ShowOpenFileDialog(app.hwnd, L"Select Video File", VideoFilters());
|
||||
ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
|
||||
} else if (cmd == "dialog-save-image") {
|
||||
std::wstring suggested = Utf8ToWide(params.value("suggestedName", std::string("output.png")));
|
||||
auto path = ShowSaveFileDialog(app.hwnd, L"Save Processed Image As",
|
||||
{ {L"PNG Image", L"*.png"}, {L"WebP Image", L"*.webp"}, {L"JPEG Image", L"*.jpg"} },
|
||||
L"png", suggested);
|
||||
ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
|
||||
} else if (cmd == "dialog-save-video") {
|
||||
std::wstring suggested = Utf8ToWide(params.value("suggestedName", std::string("output.mov")));
|
||||
auto path = ShowSaveFileDialog(app.hwnd, L"Save Output Video As",
|
||||
{ {L"QuickTime MOV", L"*.mov"}, {L"WebM Video", L"*.webm"}, {L"MP4 Video", L"*.mp4"} },
|
||||
L"mov", suggested);
|
||||
ok(path ? json{ {"path", WideToUtf8(*path)} } : json{ {"path", nullptr} });
|
||||
} else if (cmd == "dialog-pick-color") {
|
||||
std::wstring initial = Utf8ToWide(params.value("initial", std::string("#00FF00")));
|
||||
auto hex = ShowColorPickerDialog(app.hwnd, initial);
|
||||
ok(hex ? json{ {"hex", WideToUtf8(*hex)} } : json{ {"hex", nullptr} });
|
||||
} else {
|
||||
fail("Unknown native command: " + cmd);
|
||||
}
|
||||
|
||||
app.webView->PostMessageToJS(Utf8ToWide(reply.dump()));
|
||||
}
|
||||
|
||||
void OnWebMessage(AppState& app, const std::wstring& jsonStr) {
|
||||
json req;
|
||||
try {
|
||||
req = json::parse(WideToUtf8(jsonStr));
|
||||
} catch (const std::exception&) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (req.value("target", "") == "native") {
|
||||
HandleNativeCommand(app, req);
|
||||
return;
|
||||
}
|
||||
|
||||
// Everything else is a sidecar command; forward the raw JSON line as-is.
|
||||
if (app.sidecar) {
|
||||
app.sidecar->SendCommand(req.dump());
|
||||
}
|
||||
}
|
||||
|
||||
void OnSidecarEvent(AppState& app, const std::string& jsonLine) {
|
||||
// Marshal from the SidecarProcess background thread to the UI thread;
|
||||
// WebView2's COM objects are single-threaded apartment and must only
|
||||
// be touched from the thread that created them.
|
||||
auto* payload = new std::wstring(Utf8ToWide(jsonLine));
|
||||
PostMessageW(app.hwnd, WM_APP_SIDECAR_EVENT, 0, reinterpret_cast<LPARAM>(payload));
|
||||
}
|
||||
|
||||
LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
|
||||
AppState* app = g_app;
|
||||
|
||||
switch (msg) {
|
||||
case WM_SIZE:
|
||||
if (app && app->webView) app->webView->Resize();
|
||||
return 0;
|
||||
|
||||
case WM_GETMINMAXINFO: {
|
||||
auto* mmi = reinterpret_cast<MINMAXINFO*>(lParam);
|
||||
mmi->ptMinTrackSize = { 760, 560 };
|
||||
return 0;
|
||||
}
|
||||
|
||||
case WM_APP_SIDECAR_EVENT: {
|
||||
std::unique_ptr<std::wstring> payload(reinterpret_cast<std::wstring*>(lParam));
|
||||
if (app && app->webView && app->webView->IsReady()) {
|
||||
app->webView->PostMessageToJS(*payload);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
case WM_DESTROY:
|
||||
PostQuitMessage(0);
|
||||
return 0;
|
||||
}
|
||||
|
||||
return DefWindowProcW(hwnd, msg, wParam, lParam);
|
||||
}
|
||||
|
||||
void ApplyDarkTitleBar(HWND hwnd) {
|
||||
BOOL dark = TRUE;
|
||||
DwmSetWindowAttribute(hwnd, DWMWA_USE_IMMERSIVE_DARK_MODE, &dark, sizeof(dark));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int APIENTRY wWinMain(HINSTANCE hInstance, HINSTANCE, PWSTR, int nCmdShow) {
|
||||
CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
|
||||
|
||||
WNDCLASSEXW wc{};
|
||||
wc.cbSize = sizeof(wc);
|
||||
wc.style = CS_HREDRAW | CS_VREDRAW;
|
||||
wc.lpfnWndProc = WndProc;
|
||||
wc.hInstance = hInstance;
|
||||
wc.hCursor = LoadCursorW(nullptr, IDC_ARROW);
|
||||
wc.hbrBackground = CreateSolidBrush(RGB(0x11, 0x11, 0x14));
|
||||
wc.lpszClassName = kWindowClass;
|
||||
RegisterClassExW(&wc);
|
||||
|
||||
HWND hwnd = CreateWindowExW(
|
||||
0, kWindowClass, kWindowTitle,
|
||||
WS_OVERLAPPEDWINDOW,
|
||||
CW_USEDEFAULT, CW_USEDEFAULT, 1180, 820,
|
||||
nullptr, nullptr, hInstance, nullptr);
|
||||
|
||||
if (!hwnd) return 1;
|
||||
|
||||
ApplyDarkTitleBar(hwnd);
|
||||
|
||||
AppState app;
|
||||
app.hwnd = hwnd;
|
||||
g_app = &app;
|
||||
|
||||
ShowWindow(hwnd, nCmdShow);
|
||||
UpdateWindow(hwnd);
|
||||
|
||||
std::wstring repoRoot = FindRepoRoot();
|
||||
std::wstring pythonExe = ResolvePythonInterpreter(repoRoot);
|
||||
std::wstring sidecarScript = repoRoot + L"\\core\\sidecar.py";
|
||||
|
||||
app.sidecar = std::make_unique<SidecarProcess>(
|
||||
pythonExe, sidecarScript, repoRoot,
|
||||
[&app](const std::string& line) { OnSidecarEvent(app, line); });
|
||||
|
||||
if (!app.sidecar->Start()) {
|
||||
MessageBoxW(hwnd,
|
||||
L"Could not start the RemoveBG sidecar process (Python).\n"
|
||||
L"Make sure Python 3.10+ and the project dependencies are installed\n"
|
||||
L"(see README.md), then relaunch.",
|
||||
L"RemoveBG Studio", MB_ICONWARNING);
|
||||
}
|
||||
|
||||
std::wstring userDataFolder = repoRoot + L"\\native\\.webview2";
|
||||
CreateDirectoryW(userDataFolder.c_str(), nullptr);
|
||||
|
||||
#ifdef RBG_CONSOLE
|
||||
bool devTools = true;
|
||||
#else
|
||||
bool devTools = false;
|
||||
#endif
|
||||
|
||||
app.webView = std::make_unique<WebViewHost>(hwnd, userDataFolder, devTools);
|
||||
app.webView->SetMessageHandler([&app](const std::wstring& j) { OnWebMessage(app, j); });
|
||||
|
||||
std::wstring indexPath = GetExeDir() + L"\\web\\index.html";
|
||||
app.webView->Initialize([&app, indexPath]() {
|
||||
app.webView->Navigate(L"file:///" + indexPath);
|
||||
});
|
||||
|
||||
MSG msg;
|
||||
while (GetMessageW(&msg, nullptr, 0, 0)) {
|
||||
TranslateMessage(&msg);
|
||||
DispatchMessageW(&msg);
|
||||
}
|
||||
|
||||
app.sidecar.reset();
|
||||
CoUninitialize();
|
||||
return (int)msg.wParam;
|
||||
}
|
||||
Reference in New Issue
Block a user