// VIP抠图处理模块(使用 BiRefNet) // 接收 base64 图片,返回抠图后的 base64 图片 const { spawn } = require('child_process'); const path = require('path'); const fs = require('fs').promises; const fsSync = require('fs'); const os = require('os'); const http = require('http'); class MattingVIP { /** * 处理 VIP 抠图请求 * @param {http.IncomingMessage} req - HTTP 请求对象 * @param {http.ServerResponse} res - HTTP 响应对象 */ static async handleRequest(req, res) { if (req.method !== 'POST') { res.writeHead(405, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: 'Method not allowed' })); return; } let body = ''; req.on('data', (chunk) => { body += chunk.toString(); }); req.on('end', async () => { try { const data = JSON.parse(body); const { imageBase64 } = data; if (!imageBase64) { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: 'Missing required field: imageBase64' })); return; } // 调用抠图处理 const result = await this.processMatting(imageBase64); if (result.success) { res.writeHead(200, { 'Content-Type': 'application/json', 'Access-Control-Allow-Origin': '*' }); res.end(JSON.stringify({ success: true, imageData: result.imageData })); } else { res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: result.error || 'VIP抠图失败' })); } } catch (error) { console.error('[MattingVIP] 处理请求失败:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: '处理失败', details: error.message })); } }); req.on('error', (error) => { console.error('[MattingVIP] 请求错误:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: 'Request error', details: error.message })); }); } /** * 对 base64 图片进行 VIP 抠图(使用 BiRefNet) * @param {string} imageBase64 - base64 图片数据 * @returns {Promise} 处理结果 */ static async processMatting(imageBase64) { const tempDir = os.tmpdir(); const timestamp = Date.now(); const inputPath = path.join(tempDir, `matting_vip_input_${timestamp}.png`); const outputPath = path.join(tempDir, `matting_vip_output_${timestamp}.png`); try { // 清理 Base64 数据:移除可能存在的 data URL 前缀 let cleanImageData = imageBase64; if (cleanImageData.includes(',')) { const base64Index = cleanImageData.indexOf(','); cleanImageData = cleanImageData.substring(base64Index + 1); } cleanImageData = cleanImageData.trim(); // 将 base64 转换为图片文件 const imageBuffer = Buffer.from(cleanImageData, 'base64'); await fs.writeFile(inputPath, imageBuffer); console.log('[MattingVIP] 临时文件已创建:', inputPath); // 调用 Python 脚本进行抠图(使用 BiRefNet) const pythonScript = path.join(__dirname, 'python', 'birefnet-matting.py'); const result = await this.runPythonScript(pythonScript, inputPath, outputPath); if (!result.success) { return { success: false, error: result.error }; } // 读取输出图片并转换为 base64 const outputBuffer = await fs.readFile(outputPath); const outputBase64 = outputBuffer.toString('base64'); // 清理临时文件 await this.cleanupFiles([inputPath, outputPath]); console.log('[MattingVIP] ✓ VIP抠图完成'); return { success: true, imageData: outputBase64 }; } catch (error) { console.error('[MattingVIP] VIP抠图失败:', error); // 清理临时文件 await this.cleanupFiles([inputPath, outputPath]).catch(() => {}); return { success: false, error: error.message }; } } /** * 运行 Python 脚本 * @param {string} scriptPath - Python 脚本路径 * @param {string} inputPath - 输入文件路径 * @param {string} outputPath - 输出文件路径 * @returns {Promise} 执行结果 */ static async runPythonScript(scriptPath, inputPath, outputPath) { return new Promise((resolve) => { // 检查是否有虚拟环境 const venvPython = path.join(__dirname, 'python', 'venv', 'Scripts', 'python.exe'); let pythonCmd; // 优先使用虚拟环境中的 Python if (process.platform === 'win32' && fsSync.existsSync(venvPython)) { pythonCmd = venvPython; } else { pythonCmd = process.platform === 'win32' ? 'python' : 'python3'; } const python = spawn(pythonCmd, [scriptPath, inputPath, outputPath], { cwd: path.join(__dirname, 'python') }); let stdout = ''; let stderr = ''; python.stdout.on('data', (data) => { stdout += data.toString(); console.log('[MattingVIP] Python输出:', data.toString().trim()); }); python.stderr.on('data', (data) => { stderr += data.toString(); console.error('[MattingVIP] Python错误:', data.toString().trim()); }); python.on('close', (code) => { if (code === 0) { resolve({ success: true }); } else { console.error('[MattingVIP] Python 脚本执行失败:', stderr); resolve({ success: false, error: stderr || 'Python 脚本执行失败' }); } }); python.on('error', (error) => { console.error('[MattingVIP] 启动 Python 失败:', error); resolve({ success: false, error: error.message }); }); }); } /** * 清理临时文件 * @param {string[]} filePaths - 文件路径数组 */ static async cleanupFiles(filePaths) { for (const filePath of filePaths) { try { await fs.unlink(filePath); } catch (error) { // 忽略文件不存在的错误 } } } } module.exports = MattingVIP;