AnimationManager/server/matting-vip.js

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