AnimationManager/server/remove-background-base64.js

// 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<Object>} 处理结果
   */
  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<Object>} 执行结果
   */
  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;