// Base64 图片抠图处理模块 // 接收 base64 图片,返回抠图后的 base64 图片 const { spawn } = require('child_process'); const path = require('path'); const fs = require('fs').promises; const fsSync = require('fs'); const os = require('os'); class RemoveBackgroundBase64 { /** * 处理 base64 图片抠图请求 * @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.removeBackground(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 || '抠图失败' })); } } catch (error) { console.error('[RemoveBackgroundBase64] 处理请求失败:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: '处理失败', details: error.message })); } }); req.on('error', (error) => { console.error('[RemoveBackgroundBase64] 请求错误:', error); res.writeHead(500, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: 'Request error', details: error.message })); }); } /** * 对 base64 图片进行抠图 * @param {string} imageBase64 - base64 图片数据 * @returns {Promise} 处理结果 */ static async removeBackground(imageBase64) { const tempDir = os.tmpdir(); const inputPath = path.join(tempDir, `input_${Date.now()}.png`); const outputPath = path.join(tempDir, `output_${Date.now()}.png`); try { // 将 base64 转换为图片文件 const imageBuffer = Buffer.from(imageBase64, 'base64'); await fs.writeFile(inputPath, imageBuffer); console.log('[RemoveBackgroundBase64] 临时文件已创建:', inputPath); // 调用 Python 脚本进行抠图 const pythonScript = path.join(__dirname, 'python', 'rembg-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('[RemoveBackgroundBase64] ✓ 抠图完成'); return { success: true, imageData: outputBase64 }; } catch (error) { console.error('[RemoveBackgroundBase64] 抠图失败:', 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 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 pythonProcess = spawn(pythonCmd, [scriptPath, inputPath, outputPath], { cwd: path.dirname(scriptPath) }); let stdout = ''; let stderr = ''; pythonProcess.stdout.on('data', (data) => { stdout += data.toString(); }); pythonProcess.stderr.on('data', (data) => { stderr += data.toString(); }); pythonProcess.on('close', (code) => { if (code === 0) { resolve({ success: true }); } else { resolve({ success: false, error: `Python脚本执行失败,退出码: ${code}\n${stderr}` }); } }); pythonProcess.on('error', (error) => { resolve({ success: false, error: `启动Python进程失败: ${error.message}` }); }); }); } /** * 清理临时文件 * @param {string[]} filePaths - 文件路径数组 */ static async cleanupFiles(filePaths) { for (const filePath of filePaths) { try { await fs.unlink(filePath); } catch (error) { // 忽略删除失败的错误 } } } } module.exports = RemoveBackgroundBase64;