remove-background-base64.js 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165
  1. // Base64 图片抠图处理模块
  2. // 接收 base64 图片,返回抠图后的 base64 图片
  3. const { spawn } = require('child_process');
  4. const path = require('path');
  5. const fs = require('fs').promises;
  6. const os = require('os');
  7. class RemoveBackgroundBase64 {
  8. /**
  9. * 处理 base64 图片抠图请求
  10. * @param {http.IncomingMessage} req - HTTP 请求对象
  11. * @param {http.ServerResponse} res - HTTP 响应对象
  12. */
  13. static async handleRequest(req, res) {
  14. if (req.method !== 'POST') {
  15. res.writeHead(405, { 'Content-Type': 'application/json' });
  16. res.end(JSON.stringify({ error: 'Method not allowed' }));
  17. return;
  18. }
  19. let body = '';
  20. req.on('data', (chunk) => {
  21. body += chunk.toString();
  22. });
  23. req.on('end', async () => {
  24. try {
  25. const data = JSON.parse(body);
  26. const { imageBase64 } = data;
  27. if (!imageBase64) {
  28. res.writeHead(400, { 'Content-Type': 'application/json' });
  29. res.end(JSON.stringify({ error: 'Missing required field: imageBase64' }));
  30. return;
  31. }
  32. // 调用抠图处理
  33. const result = await this.removeBackground(imageBase64);
  34. if (result.success) {
  35. res.writeHead(200, {
  36. 'Content-Type': 'application/json',
  37. 'Access-Control-Allow-Origin': '*'
  38. });
  39. res.end(JSON.stringify({
  40. success: true,
  41. imageData: result.imageData
  42. }));
  43. } else {
  44. res.writeHead(500, { 'Content-Type': 'application/json' });
  45. res.end(JSON.stringify({ error: result.error || '抠图失败' }));
  46. }
  47. } catch (error) {
  48. console.error('[RemoveBackgroundBase64] 处理请求失败:', error);
  49. res.writeHead(500, { 'Content-Type': 'application/json' });
  50. res.end(JSON.stringify({ error: '处理失败', details: error.message }));
  51. }
  52. });
  53. req.on('error', (error) => {
  54. console.error('[RemoveBackgroundBase64] 请求错误:', error);
  55. res.writeHead(500, { 'Content-Type': 'application/json' });
  56. res.end(JSON.stringify({ error: 'Request error', details: error.message }));
  57. });
  58. }
  59. /**
  60. * 对 base64 图片进行抠图
  61. * @param {string} imageBase64 - base64 图片数据
  62. * @returns {Promise<Object>} 处理结果
  63. */
  64. static async removeBackground(imageBase64) {
  65. const tempDir = os.tmpdir();
  66. const inputPath = path.join(tempDir, `input_${Date.now()}.png`);
  67. const outputPath = path.join(tempDir, `output_${Date.now()}.png`);
  68. try {
  69. // 将 base64 转换为图片文件
  70. const imageBuffer = Buffer.from(imageBase64, 'base64');
  71. await fs.writeFile(inputPath, imageBuffer);
  72. console.log('[RemoveBackgroundBase64] 临时文件已创建:', inputPath);
  73. // 调用 Python 脚本进行抠图
  74. const pythonScript = path.join(__dirname, 'python', 'remove-background-base64.py');
  75. const result = await this.runPythonScript(pythonScript, inputPath, outputPath);
  76. if (!result.success) {
  77. return { success: false, error: result.error };
  78. }
  79. // 读取输出图片并转换为 base64
  80. const outputBuffer = await fs.readFile(outputPath);
  81. const outputBase64 = outputBuffer.toString('base64');
  82. // 清理临时文件
  83. await this.cleanupFiles([inputPath, outputPath]);
  84. console.log('[RemoveBackgroundBase64] ✓ 抠图完成');
  85. return { success: true, imageData: outputBase64 };
  86. } catch (error) {
  87. console.error('[RemoveBackgroundBase64] 抠图失败:', error);
  88. // 清理临时文件
  89. await this.cleanupFiles([inputPath, outputPath]).catch(() => {});
  90. return { success: false, error: error.message };
  91. }
  92. }
  93. /**
  94. * 运行 Python 脚本
  95. * @param {string} scriptPath - Python 脚本路径
  96. * @param {string} inputPath - 输入文件路径
  97. * @param {string} outputPath - 输出文件路径
  98. * @returns {Promise<Object>} 执行结果
  99. */
  100. static runPythonScript(scriptPath, inputPath, outputPath) {
  101. return new Promise((resolve) => {
  102. const pythonProcess = spawn('python', [scriptPath, inputPath, outputPath], {
  103. cwd: path.dirname(scriptPath)
  104. });
  105. let stdout = '';
  106. let stderr = '';
  107. pythonProcess.stdout.on('data', (data) => {
  108. stdout += data.toString();
  109. });
  110. pythonProcess.stderr.on('data', (data) => {
  111. stderr += data.toString();
  112. });
  113. pythonProcess.on('close', (code) => {
  114. if (code === 0) {
  115. resolve({ success: true });
  116. } else {
  117. resolve({ success: false, error: `Python脚本执行失败,退出码: ${code}\n${stderr}` });
  118. }
  119. });
  120. pythonProcess.on('error', (error) => {
  121. resolve({ success: false, error: `启动Python进程失败: ${error.message}` });
  122. });
  123. });
  124. }
  125. /**
  126. * 清理临时文件
  127. * @param {string[]} filePaths - 文件路径数组
  128. */
  129. static async cleanupFiles(filePaths) {
  130. for (const filePath of filePaths) {
  131. try {
  132. await fs.unlink(filePath);
  133. } catch (error) {
  134. // 忽略删除失败的错误
  135. }
  136. }
  137. }
  138. }
  139. module.exports = RemoveBackgroundBase64;