feat(demo): 新增图片生成与保存功能,优化图片提取逻辑

This commit is contained in:
WJZ_P
2026-03-18 02:23:43 +08:00
parent d7f09ff8f7
commit 464327f24b
2 changed files with 130 additions and 39 deletions

View File

@@ -15,10 +15,12 @@
* 所有配置项见 .env可直接编辑或通过命令行设环境变量。 * 所有配置项见 .env可直接编辑或通过命令行设环境变量。
*/ */
import { execSync } from 'node:child_process'; import { execSync } from 'node:child_process';
import { platform } from 'node:os'; import { platform, homedir } from 'node:os';
import { writeFileSync, mkdirSync, existsSync } from 'node:fs';
import { join } from 'node:path';
import { createGeminiSession, disconnect } from './index.js'; import { createGeminiSession, disconnect } from './index.js';
const prompt = 'Hello Gemini!'; const prompt = '你好呀~请给我画一张植物大战僵尸的漫画,左侧是豌豆射手,右侧是僵尸。僵尸面对豌豆射手的攻击仓皇逃窜!';
// ── Demo 专用:杀掉所有 Chromium 系浏览器进程 ── // ── Demo 专用:杀掉所有 Chromium 系浏览器进程 ──
function killAllBrowserProcesses() { function killAllBrowserProcesses() {
@@ -111,16 +113,74 @@ async function main() {
const probe = await ops.probe(); const probe = await ops.probe();
console.log('probe:', JSON.stringify(probe, null, 2)); console.log('probe:', JSON.stringify(probe, null, 2));
// 3. 发送一句话 // 3. 确保使用 Pro 模型
console.log('\n[3] 发送提示词...'); console.log('\n[3] 检查模型...');
if (probe.currentModel.toLowerCase() === 'pro') {
console.log('[3] 当前已是 Pro 模型,跳过');
} else {
console.log(`[3] 当前模型: ${probe.currentModel || '未知'},切换到 Pro...`);
const switchResult = await ops.ensureModelPro();
if (switchResult.ok) {
console.log(`[3] 已切换到 Pro之前: ${switchResult.previousModel || '未知'}`);
} else {
console.warn(`[3] 切换 Pro 失败: ${switchResult.error},继续使用当前模型`);
}
}
// 4. 发送一句话
console.log('\n[4] 发送提示词...');
const result = await ops.sendAndWait(prompt, { const result = await ops.sendAndWait(prompt, {
timeout: 60_000, timeout: 120_000,
onPoll(poll) { onPoll(poll) {
console.log(` polling... status=${poll.status}`); console.log(` polling... status=${poll.status}`);
}, },
}); });
console.log('result:', JSON.stringify(result, null, 2)); console.log('result:', JSON.stringify(result, null, 2));
// 5. 获取最新图片并保存到本地
if (result.ok) {
console.log('\n[5] 查找最新生成的图片...');
await sleep(2000); // 等待图片渲染完毕
const imgInfo = await ops.getLatestImage();
console.log('imgInfo:', JSON.stringify(imgInfo, null, 2));
if (imgInfo.ok && imgInfo.src) {
console.log(`[5] 找到图片 (${imgInfo.width}x${imgInfo.height}, isNew=${imgInfo.isNew})`);
// 提取 base64 数据
console.log(`[5] 提取图片数据 (src=${imgInfo.src})...`);
const b64Result = await ops.extractImageBase64(imgInfo.src);
if (b64Result.ok && b64Result.dataUrl) {
// dataUrl 格式: data:image/png;base64,iVBOR...
const matches = b64Result.dataUrl.match(/^data:image\/(\w+);base64,(.+)$/);
if (matches) {
const ext = matches[1] === 'jpeg' ? 'jpg' : matches[1];
const base64Data = matches[2];
const buffer = Buffer.from(base64Data, 'base64');
// 保存到 ~/gemini-skill-output/
const outputDir = join(homedir(), 'gemini-skill-output');
if (!existsSync(outputDir)) {
mkdirSync(outputDir, { recursive: true });
}
const filename = `gemini_${Date.now()}.${ext}`;
const filepath = join(outputDir, filename);
writeFileSync(filepath, buffer);
console.log(`[5] ✅ 图片已保存: ${filepath} (${(buffer.length / 1024).toFixed(1)} KB, method=${b64Result.method})`);
} else {
console.warn('[5] ⚠ dataUrl 格式无法解析');
}
} else {
console.warn(`[5] ⚠ 提取图片数据失败: ${b64Result.error || 'unknown'}`);
}
} else {
console.log('[5] 未找到图片(可能本次回答不含图片)');
}
}
} catch (err) { } catch (err) {
console.error('Error:', err); console.error('Error:', err);
} }

View File

@@ -465,19 +465,29 @@ export function createOps(page) {
}, },
/** /**
* 提取最新图片的 Base64 数据Canvas 优先fetch 兜底) * 提取指定图片的 Base64 数据Canvas 优先fetch 兜底)
*
* @param {string} url - 目标图片的 src URL
* @returns {Promise<{ok: boolean, dataUrl?: string, width?: number, height?: number, method?: 'canvas'|'fetch', error?: string}>}
*/ */
async extractImageBase64() { async extractImageBase64(url) {
return op.query(() => { if (!url) {
console.warn('[extractImageBase64] ❌ 未提供 url 参数');
return { ok: false, error: 'missing_url' };
}
console.log(`[extractImageBase64] 🔍 开始提取, url=${url.slice(0, 120)}...`);
const canvasResult = await op.query((targetUrl) => {
// ── 在页面中根据 url 查找匹配的 img 元素 ──
const imgs = [...document.querySelectorAll('img.image.loaded')]; const imgs = [...document.querySelectorAll('img.image.loaded')];
if (!imgs.length) { const img = imgs.find(i => i.src === targetUrl);
return { ok: false, error: 'no_loaded_images' }; if (!img) {
return { ok: false, error: 'img_not_found_by_url', searched: imgs.length };
} }
const img = imgs[imgs.length - 1];
const w = img.naturalWidth || img.width; const w = img.naturalWidth || img.width;
const h = img.naturalHeight || img.height; const h = img.naturalHeight || img.height;
// 尝试 Canvas 同步提取 // ── 尝试 Canvas 同步提取 ──
try { try {
const canvas = document.createElement('canvas'); const canvas = document.createElement('canvas');
canvas.width = w; canvas.width = w;
@@ -485,31 +495,52 @@ export function createOps(page) {
canvas.getContext('2d').drawImage(img, 0, 0); canvas.getContext('2d').drawImage(img, 0, 0);
const dataUrl = canvas.toDataURL('image/png'); const dataUrl = canvas.toDataURL('image/png');
return { ok: true, dataUrl, width: w, height: h, method: 'canvas' }; return { ok: true, dataUrl, width: w, height: h, method: 'canvas' };
} catch { /* canvas tainted, fallback */ } } catch (e) {
// canvas tainted跨域图片记录原因后降级
return { ok: false, needFetch: true, src: img.src, width: w, height: h, canvasError: e.message || String(e) };
}
}, url);
// 标记需要 fetch fallback if (canvasResult.ok) {
return { ok: false, needFetch: true, src: img.src, width: w, height: h }; console.log(`[extractImageBase64] ✅ Canvas 提取成功 (${canvasResult.width}x${canvasResult.height})`);
}).then(async (result) => { return canvasResult;
if (result.ok || !result.needFetch) return result; }
// Fetch fallback: 在页面上下文中异步执行 if (!canvasResult.needFetch) {
return page.evaluate(async (src, w, h) => { // img 元素都没找到,直接返回失败
try { console.warn(`[extractImageBase64] ❌ 页面中未找到匹配的 img 元素 (已扫描 ${canvasResult.searched || 0} 张)`);
const r = await fetch(src); return canvasResult;
if (!r.ok) throw new Error(`fetch_status_${r.status}`); }
const blob = await r.blob();
return await new Promise((resolve) => { // ── Fetch 降级Canvas 被跨域污染,改用 fetch 读取二进制 ──
const reader = new FileReader(); console.log(`[extractImageBase64] ⚠ Canvas 被污染 (${canvasResult.canvasError}),降级为 fetch...`);
reader.onloadend = () => resolve({
ok: true, dataUrl: reader.result, width: w, height: h, method: 'fetch', const fetchResult = await page.evaluate(async (src, w, h) => {
}); try {
reader.readAsDataURL(blob); const r = await fetch(src);
if (!r.ok) return { ok: false, error: `fetch_status_${r.status}` };
const blob = await r.blob();
return await new Promise((resolve) => {
const reader = new FileReader();
reader.onloadend = () => resolve({
ok: true, dataUrl: reader.result, width: w, height: h, method: 'fetch',
}); });
} catch (err) { reader.onerror = () => resolve({
return { ok: false, error: 'extract_failed', detail: err.message || String(err) }; ok: false, error: 'filereader_error',
} });
}, result.src, result.width, result.height); reader.readAsDataURL(blob);
}); });
} catch (err) {
return { ok: false, error: 'fetch_failed', detail: err.message || String(err) };
}
}, canvasResult.src, canvasResult.width, canvasResult.height);
if (fetchResult.ok) {
console.log(`[extractImageBase64] ✅ Fetch 提取成功 (${fetchResult.width}x${fetchResult.height})`);
} else {
console.warn(`[extractImageBase64] ❌ Fetch 提取失败: ${fetchResult.error}${fetchResult.detail ? ' — ' + fetchResult.detail : ''}`);
}
return fetchResult;
}, },
/** /**
@@ -621,12 +652,12 @@ export function createOps(page) {
await sleep(2000); await sleep(2000);
// 4. 获取图片 // 4. 获取图片
const imgInfo = await this.getLatestImage(); let imgInfo = await this.getLatestImage();
if (!imgInfo.ok) { if (!imgInfo.ok) {
await sleep(3000); await sleep(3000);
const retry = await this.getLatestImage(); imgInfo = await this.getLatestImage();
if (!retry.ok) { if (!imgInfo.ok) {
return { ok: false, error: 'no_image_found', elapsed: waitResult.elapsed, imgInfo: retry }; return { ok: false, error: 'no_image_found', elapsed: waitResult.elapsed, imgInfo };
} }
} }
@@ -635,7 +666,7 @@ export function createOps(page) {
const dlResult = await this.downloadLatestImage(); const dlResult = await this.downloadLatestImage();
return { ok: dlResult.ok, method: 'download', elapsed: waitResult.elapsed, ...dlResult }; return { ok: dlResult.ok, method: 'download', elapsed: waitResult.elapsed, ...dlResult };
} else { } else {
const b64Result = await this.extractImageBase64(); const b64Result = await this.extractImageBase64(imgInfo.src);
return { ok: b64Result.ok, method: b64Result.method, elapsed: waitResult.elapsed, ...b64Result }; return { ok: b64Result.ok, method: b64Result.method, elapsed: waitResult.elapsed, ...b64Result };
} }
}, },