GPU_GUARD_MONOREPO/packages/backend/skills/video-frame-extractor/scripts/index.cjs

#!/usr/bin/env node
'use strict';

const fs = require('node:fs');
const os = require('node:os');
const path = require('node:path');
const crypto = require('node:crypto');
const { spawn } = require('node:child_process');

function emit(stage, message, status = 'running', extra = {}) {
  process.stderr.write(JSON.stringify({
    type: 'process_event',
    stage,
    message,
    status,
    timestamp: new Date().toISOString(),
    ...extra,
  }) + '\n');
}

function readInput() {
  const raw = String(process.argv[2] || process.env.SKILL_INPUT || process.env.AIFLOW_SKILL_INPUT || fs.readFileSync(0, 'utf8')).trim();
  return raw ? JSON.parse(raw) : {};
}

function clampNumber(value, min, max, fallback) {
  const number = Number(value);
  if (!Number.isFinite(number)) return fallback;
  return Math.max(min, Math.min(max, number));
}

function safeTaskId(value) {
  const raw = String(value || '').trim();
  if (/^[a-zA-Z0-9_-]{1,80}$/.test(raw)) return raw;
  return `task_${Date.now().toString(36)}_${crypto.randomBytes(4).toString('hex')}`;
}

function workspaceRoot(env) {
  if (env.VEHICLE_SCRATCH_WORKSPACE_ROOT) return path.resolve(env.VEHICLE_SCRATCH_WORKSPACE_ROOT);
  if (env.RZYX_AI_WORKSPACE_ROOT) {
    const root = path.resolve(env.RZYX_AI_WORKSPACE_ROOT);
    return path.basename(root) === 'vehicle-scratch-inspection' ? root : path.join(root, 'vehicle-scratch-inspection');
  }
  if (env.RZYX_AI_DATA_DIR) return path.join(path.resolve(env.RZYX_AI_DATA_DIR), 'workspace', 'vehicle-scratch-inspection');
  return path.join(os.tmpdir(), 'vehicle-scratch-inspection');
}

function createWorkspace(taskId, env) {
  const root = workspaceRoot(env);
  const id = safeTaskId(taskId);
  const workspacePath = path.join(root, id);
  for (const dir of ['', 'source', 'frames', 'marked_frames', 'best_frames', 'report']) {
    fs.mkdirSync(path.join(workspacePath, dir), { recursive: true });
  }
  return { taskId: id, workspaceRoot: root, workspacePath };
}

function writeJson(filePath, value) {
  fs.mkdirSync(path.dirname(filePath), { recursive: true });
  fs.writeFileSync(filePath, `${JSON.stringify(value, null, 2)}\n`, 'utf8');
}

function runProcess(command, args, options = {}) {
  return new Promise((resolve, reject) => {
    const child = spawn(command, args, { stdio: ['ignore', 'pipe', 'pipe'], ...options });
    let stdout = '';
    let stderr = '';
    child.stdout.on('data', chunk => { stdout += chunk.toString(); });
    child.stderr.on('data', chunk => { stderr += chunk.toString(); });
    child.on('error', reject);
    child.on('close', code => {
      if (code === 0) resolve({ stdout, stderr });
      else reject(new Error(`${command} exited with ${code}: ${stderr.trim().slice(-1200)}`));
    });
  });
}

async function downloadVideo(url, targetPath) {
  emit('download', '下载远程视频到workspace');
  const response = await fetch(url);
  if (!response.ok) throw new Error(`视频下载失败: HTTP ${response.status}`);
  const buffer = Buffer.from(await response.arrayBuffer());
  fs.writeFileSync(targetPath, buffer);
  return targetPath;
}

async function resolveVideo(videoUrl, workspace, env) {
  const value = String(videoUrl || '').trim();
  if (!value) throw new Error('videoUrl不能为空');
  if (fs.existsSync(value)) return path.resolve(value);

  const normalized = value.replace(/\\/g, '/');
  if (normalized.startsWith('/upload/')) {
    const uploadRoot = env.RZYX_AI_UPLOAD_ROOT || (env.RZYX_AI_DATA_DIR ? path.join(env.RZYX_AI_DATA_DIR, 'uploads') : '');
    if (!uploadRoot) throw new Error('缺少RZYX_AI_UPLOAD_ROOT，无法解析/upload路径');
    const relative = normalized.slice('/upload/'.length);
    const resolved = path.resolve(uploadRoot, relative);
    const root = path.resolve(uploadRoot);
    if (!resolved.startsWith(root + path.sep)) throw new Error('upload路径非法');
    if (!fs.existsSync(resolved)) throw new Error(`视频文件不存在: ${resolved}`);
    return resolved;
  }

  if (/^https?:\/\//i.test(value)) {
    const url = new URL(value);
    if (['localhost', '127.0.0.1'].includes(url.hostname)) {
      try {
        return await resolveVideo(decodeURIComponent(url.pathname), workspace, env);
      } catch {
        // Fall through to HTTP download for local static servers that do not map to files.
      }
    }
    const ext = path.extname(url.pathname) || '.mp4';
    return downloadVideo(value, path.join(workspace.workspacePath, 'source', `input${ext}`));
  }

  throw new Error(`视频文件不存在: ${value}`);
}

function parseRatio(value) {
  const text = String(value || '');
  const match = text.match(/^(\d+(?:\.\d+)?)\/(\d+(?:\.\d+)?)$/);
  if (match) {
    const a = Number(match[1]);
    const b = Number(match[2]);
    return b ? a / b : 0;
  }
  const n = Number(text);
  return Number.isFinite(n) ? n : 0;
}

function ffmpegCommand(env) {
  if (env.FFMPEG_PATH) return env.FFMPEG_PATH;
  try { return require('@ffmpeg-installer/ffmpeg').path; } catch {}
  return 'ffmpeg';
}

function ffprobeCommand(env) {
  if (env.FFPROBE_PATH) return env.FFPROBE_PATH;
  try { return require('@ffprobe-installer/ffprobe').path; } catch {}
  return 'ffprobe';
}

async function probeVideo(videoPath, env) {
  const { stdout } = await runProcess(ffprobeCommand(env), [
    '-v', 'error',
    '-print_format', 'json',
    '-show_format',
    '-show_streams',
    videoPath,
  ]);
  const meta = JSON.parse(stdout || '{}');
  const stream = (meta.streams || []).find(item => item.codec_type === 'video') || {};
  const duration = Number(meta.format?.duration || stream.duration || 0);
  const videoFps = parseRatio(stream.avg_frame_rate || stream.r_frame_rate);
  const totalFramesRaw = Number(stream.nb_frames);
  const totalFrames = Number.isFinite(totalFramesRaw) && totalFramesRaw > 0
    ? totalFramesRaw
    : Math.round(duration * (videoFps || 0));
  return {
    totalFrames,
    videoFps: Number((videoFps || 0).toFixed(3)),
    duration: Number((duration || 0).toFixed(3)),
    resolution: `${stream.width || 0}x${stream.height || 0}`,
  };
}

async function extractFrames({ videoPath, workspace, fps, quality, env }) {
  const framesDir = path.join(workspace.workspacePath, 'frames');
  for (const file of fs.readdirSync(framesDir)) {
    if (/\.jpe?g$/i.test(file)) fs.rmSync(path.join(framesDir, file), { force: true });
  }
  const qscale = Math.max(2, Math.min(31, Math.round(31 - (quality / 100) * 29)));
  emit('extract_frames', '开始按指定fps抽帧', 'running', { fps, quality });
  await runProcess(ffmpegCommand(env), [
    '-hide_banner',
    '-y',
    '-i', videoPath,
    '-vf', `fps=${fps}`,
    '-q:v', String(qscale),
    path.join(framesDir, 'frame_%06d.jpg'),
  ]);
  const files = fs.readdirSync(framesDir).filter(file => /\.jpe?g$/i.test(file)).sort();
  return files.map((file, index) => ({
    index: index + 1,
    timestamp: Number((index / fps).toFixed(3)),
    fileName: file,
    path: path.join(framesDir, file),
    relativePath: `frames/${file}`,
  }));
}

async function run(input, env = process.env) {
  if (!input || typeof input !== 'object') throw new Error('输入必须是JSON对象');
  const fps = clampNumber(input.fps, 0.2, 10, 5);
  const quality = clampNumber(input.quality, 1, 100, 90);
  const workspace = createWorkspace(input.taskId, env);
  emit('workspace', 'workspace已创建', 'running', { taskId: workspace.taskId, workspacePath: workspace.workspacePath });

  const videoPath = await resolveVideo(input.videoUrl, workspace, env);
  const baseInfo = await probeVideo(videoPath, env);
  const frames = await extractFrames({ videoPath, workspace, fps, quality, env });
  if (frames.length === 0) throw new Error('视频无法抽帧，请检查视频编码或fps参数');

  const videoInfo = {
    ...baseInfo,
    extractedFrames: frames.length,
    extractFps: fps,
  };
  const result = {
    success: true,
    taskId: workspace.taskId,
    workspacePath: workspace.workspacePath,
    frameCount: frames.length,
    videoInfo,
  };
  writeJson(path.join(workspace.workspacePath, 'video_info.json'), {
    taskId: workspace.taskId,
    workspacePath: workspace.workspacePath,
    videoUrl: input.videoUrl,
    videoPath,
    frames,
    videoInfo,
  });
  emit('extract_frames', '抽帧完成并写入video_info.json', 'completed', { frameCount: frames.length });
  return result;
}

(async () => {
  try {
    process.stdout.write(JSON.stringify(await run(readInput())));
  } catch (err) {
    process.stdout.write(JSON.stringify({
      success: false,
      error: err && err.message ? err.message : String(err),
    }));
  }
})();