228 lines
9.6 KiB
TypeScript
228 lines
9.6 KiB
TypeScript
import { callLmstudio } from '../../lib/lmstudio';
|
|
import { logger } from '../../lib/logger';
|
|
import * as fs from 'fs/promises';
|
|
import dotenv from 'dotenv';
|
|
import path from 'path';
|
|
import puppeteer from 'puppeteer';
|
|
|
|
dotenv.config();
|
|
|
|
const SCROLL_SEARCH = 3; // scroll times on search results
|
|
const SCROLL_PIN = 3; // scroll times on pin page
|
|
const PINS_TO_COLLECT = 5;
|
|
|
|
// Hard-coded user prompt
|
|
const HARDCODED_USER_PROMPT = process.env.HARDCODED_USER_PROMPT || `
|
|
Generate 20 keywords for photos of a ghotst or monster from all over the world. "Cute Japanese yokai" is mandatory, also add "Realistic photo cute" keyword to all genearated keywords first.
|
|
Example output : ["Cute Japanese yokai","Realistic photo Cute ghost","Realistic photo cute monster","Realistic photo cute haloween monster","Realistic photo cute haloween ghost"... and 20 items in array]
|
|
`;
|
|
|
|
async function getPinUrlsFromPinterest(keyword: string, scrollCount = SCROLL_SEARCH, limit = PINS_TO_COLLECT): Promise<string[]> {
|
|
const browser = await puppeteer.launch({ headless: true });
|
|
const page = await browser.newPage();
|
|
await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36');
|
|
await page.setViewport({ width: 1920, height: 1080 });
|
|
try {
|
|
const searchUrl = `https://www.pinterest.com/search/pins/?q=${encodeURIComponent(keyword)}`;
|
|
await page.goto(searchUrl, { waitUntil: 'networkidle2' });
|
|
|
|
let pinLinks = new Set<string>();
|
|
for (let i = 0; i < scrollCount; i++) {
|
|
const linksBefore = pinLinks.size;
|
|
const newLinks = await page.$$eval('a', (anchors) =>
|
|
anchors.map((a) => a.href).filter((href) => href.includes('/pin/'))
|
|
);
|
|
newLinks.forEach(link => pinLinks.add(link));
|
|
|
|
if (pinLinks.size >= limit) {
|
|
break;
|
|
}
|
|
|
|
await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');
|
|
await new Promise(r => setTimeout(r, 500 + Math.random() * 1000));
|
|
|
|
if (pinLinks.size === linksBefore) {
|
|
// If no new pins are loaded, stop scrolling
|
|
logger.info(`No new pins loaded for "${keyword}", stopping scroll.`);
|
|
break;
|
|
}
|
|
}
|
|
return Array.from(pinLinks).slice(0, limit);
|
|
} catch (error) {
|
|
logger.error(`Error while getting pin URLs from Pinterest for keyword "${keyword}":`, error);
|
|
return [];
|
|
} finally {
|
|
await browser.close();
|
|
}
|
|
}
|
|
|
|
async function downloadImagesFromPin(pinUrl: string, scrollTimes = SCROLL_PIN): Promise<string[]> {
|
|
const browser = await puppeteer.launch({ headless: true });
|
|
const page = await browser.newPage();
|
|
await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36');
|
|
await page.setViewport({ width: 1920, height: 1080 });
|
|
try {
|
|
await page.goto(pinUrl, { waitUntil: 'networkidle2', timeout: 30000 });
|
|
for (let i = 0; i < scrollTimes; i++) {
|
|
await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');
|
|
await new Promise((r) => setTimeout(r, 700 + Math.random() * 800));
|
|
}
|
|
const imgs: string[] = await page.$$eval('img', imgs => {
|
|
const urls: string[] = imgs.map(img => {
|
|
const srcset = (img as HTMLImageElement).getAttribute('srcset') || '';
|
|
if (!srcset) {
|
|
return ''; // Ignore images without srcset
|
|
}
|
|
|
|
const parts = srcset.split(',').map(p => p.trim());
|
|
for (const part of parts) {
|
|
const match = part.match(/^(\S+)\s+4x$/);
|
|
if (match && match[1]) {
|
|
return match[1]; // Found the 4x version, return it
|
|
}
|
|
}
|
|
|
|
return ''; // No 4x version found for this image
|
|
}).filter(s => !!s && s.includes('pinimg')); // Filter out empty strings and non-pinterest images
|
|
return [...new Set(urls)]; // Return unique URLs
|
|
});
|
|
|
|
if (!imgs || imgs.length === 0) {
|
|
logger.warn(`No high-res images found on pin ${pinUrl}`);
|
|
return [];
|
|
}
|
|
|
|
const outDir = path.join(process.cwd(), 'download');
|
|
await fs.mkdir(outDir, { recursive: true });
|
|
const results: string[] = [];
|
|
for (let i = 0; i < imgs.length; i++) {
|
|
const src = imgs[i];
|
|
try {
|
|
const imgPage = await browser.newPage();
|
|
const resp = await imgPage.goto(src, { timeout: 30000, waitUntil: 'load' });
|
|
if (!resp) { await imgPage.close(); continue; }
|
|
const buffer = await resp.buffer();
|
|
const pinId = pinUrl.split('/').filter(Boolean).pop() || `pin_${Date.now()}`;
|
|
const timestamp = Date.now();
|
|
const outPath = path.join(outDir, `${pinId}_${timestamp}_${i}.png`);
|
|
await fs.writeFile(outPath, buffer);
|
|
results.push(outPath);
|
|
await imgPage.close();
|
|
} catch (err) {
|
|
logger.error(`Failed to download image ${src} from ${pinUrl}:`, err);
|
|
}
|
|
}
|
|
return results;
|
|
} catch (err) {
|
|
logger.error(`Failed to download images from ${pinUrl}:`, err);
|
|
return [];
|
|
} finally {
|
|
await browser.close();
|
|
}
|
|
}
|
|
|
|
|
|
// Re-usable helper to extract JSON embedded in text
|
|
function extractJsonFromText(text: string): any | null {
|
|
if (!text || typeof text !== 'string') return null;
|
|
const fenced = text.match(/```(?:json)?\s*([\s\S]*?)\s*```/i);
|
|
if (fenced && fenced[1]) {
|
|
try { return JSON.parse(fenced[1].trim()); } catch (e) { /* fall through */ }
|
|
}
|
|
const brace = text.match(/\{[\s\S]*\}|\[[\s\S]*\]/);
|
|
if (brace && brace[0]) {
|
|
try { return JSON.parse(brace[0]); } catch (e) { return null; }
|
|
}
|
|
// Attempt line-separated keywords fallback
|
|
const lines = text.split(/\r?\n/).map((l: string) => l.trim()).filter(Boolean);
|
|
if (lines.length > 1) return lines;
|
|
return null;
|
|
}
|
|
|
|
async function extractKeywordsFromPromptWithLmstudio(prompt: string, count = 5): Promise<string[]> {
|
|
const instruction = `You are given a short instruction describing the type of content to search for.
|
|
Return exactly a JSON array of ${count} short keyword phrases suitable for searching Pinterest. `;
|
|
|
|
try {
|
|
const res = await callLmstudio(`${instruction}\n\nInstruction: ${prompt}`);
|
|
if (!res) {
|
|
logger.warn('callLmstudio returned empty response for keyword extraction.');
|
|
return [];
|
|
}
|
|
|
|
let parsed: any;
|
|
if (typeof res === 'object' && res.text) {
|
|
parsed = extractJsonFromText(res.text);
|
|
} else if (typeof res === 'string') {
|
|
parsed = extractJsonFromText(res);
|
|
} else if (typeof res === 'object') {
|
|
parsed = res;
|
|
}
|
|
|
|
if (Array.isArray(parsed)) {
|
|
return parsed.map(String).slice(0, count);
|
|
}
|
|
|
|
if (typeof parsed === 'object' && parsed !== null) {
|
|
const maybe = parsed.keywords || parsed.list || parsed.items || parsed.keywords_list;
|
|
if (Array.isArray(maybe)) return maybe.map(String).slice(0, count);
|
|
}
|
|
|
|
const text = typeof res === 'string' ? res : (res && res.text) || JSON.stringify(res);
|
|
const lines = text.split(/\r?\n/).map((l: string) => l.replace(/^\d+[\).\s-]*/, '').trim()).filter(Boolean);
|
|
if (lines.length >= 1) {
|
|
return lines.slice(0, count);
|
|
}
|
|
|
|
logger.warn(`Could not parse keywords from LM Studio response: ${JSON.stringify(res)}`);
|
|
return [];
|
|
|
|
} catch (error) {
|
|
logger.error('Error during keyword extraction with callLmstudio:', error);
|
|
return [];
|
|
}
|
|
}
|
|
|
|
|
|
(async () => {
|
|
logger.info(`Starting photo download process with prompt: "${HARDCODED_USER_PROMPT}"`);
|
|
|
|
// 1. Extract keywords from the hardcoded prompt
|
|
const keywords = await extractKeywordsFromPromptWithLmstudio(HARDCODED_USER_PROMPT, 20); // Using 5 keywords to get a good variety
|
|
if (!keywords || keywords.length === 0) {
|
|
logger.error("Could not extract keywords from prompt. Exiting.");
|
|
return;
|
|
}
|
|
logger.info(`Extracted keywords: ${keywords.join(', ')}`);
|
|
|
|
// 2. Search Pinterest for each keyword and collect pin URLs
|
|
let allPinUrls = new Set<string>();
|
|
for (const keyword of keywords) {
|
|
logger.info(`Searching Pinterest for keyword: "${keyword}"`);
|
|
const pinUrls = await getPinUrlsFromPinterest(keyword, SCROLL_SEARCH, PINS_TO_COLLECT);
|
|
pinUrls.forEach(url => allPinUrls.add(url));
|
|
}
|
|
|
|
const finalPinUrls = Array.from(allPinUrls);
|
|
logger.info(`Collected ${finalPinUrls.length} unique pin URLs to process.`);
|
|
|
|
// 3. Go through each pin URL, scroll, and download all photos
|
|
let totalDownloads = 0;
|
|
for (const pinUrl of finalPinUrls) {
|
|
try {
|
|
logger.info(`Processing pin: ${pinUrl}`);
|
|
const downloadedPaths = await downloadImagesFromPin(pinUrl, SCROLL_PIN);
|
|
if (downloadedPaths.length > 0) {
|
|
logger.info(`Successfully downloaded ${downloadedPaths.length} images from ${pinUrl}`);
|
|
totalDownloads += downloadedPaths.length;
|
|
} else {
|
|
logger.warn(`No images were downloaded from ${pinUrl}`);
|
|
}
|
|
} catch (error) {
|
|
logger.error(`An error occurred while processing pin ${pinUrl}:`, error);
|
|
}
|
|
}
|
|
|
|
logger.info(`Photo download process finished. Total images downloaded: ${totalDownloads}`);
|
|
})();
|