239 lines
8.4 KiB
JavaScript
239 lines
8.4 KiB
JavaScript
#!/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),
|
||
}));
|
||
}
|
||
})();
|