image-dataset/dist/collect.js

Tool to build image dataset: collect, classify, review

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.collectByKeyword = collectByKeyword;
exports.main = main;
const progress_cli_1 = require("@beenotung/tslib/progress-cli");
const crypto_1 = require("crypto");
const path_1 = require("path");
const promises_1 = require("fs/promises");
const better_sqlite3_proxy_1 = require("better-sqlite3-proxy");
const proxy_1 = require("./proxy");
const config_1 = require("./config");
const browser_1 = require("./browser");
let cli = new progress_cli_1.ProgressCli();
async function collectByKeyword(keyword, options = {}) {
    let cli_prefix = options.cli_prefix || '';
    cli.update(`${cli_prefix}searching "${keyword}"...`);
    let dir = (0, path_1.join)(config_1.config.downloadedRootDir, keyword);
    await (0, promises_1.mkdir)(dir, { recursive: true });
    let keyword_id = (0, better_sqlite3_proxy_1.seedRow)(proxy_1.proxy.keyword, { keyword });
    let page = await (0, browser_1.getPage)();
    await page.goto('https://images.google.com', {
        waitUntil: 'domcontentloaded',
    });
    await page.fill('form textarea[name="q"]', keyword);
    await page.evaluate(() => {
        let form = document.querySelector('form[role="search"]');
        form.submit();
    });
    await page.waitForURL(/^https:\/\/www\.google\.com\/search/, {
        waitUntil: 'domcontentloaded',
    });
    async function collectImages() {
        let images = await page.evaluate(async () => {
            let items = document.querySelectorAll('[data-lpage]');
            let images = [];
            for (let item of items) {
                let page_url = item.dataset.lpage;
                let img = item.querySelector('img');
                let loading_gif = 'data:image/gif;base64,R0lGODlhAQABAIAAAP///////yH5BAEKAAEALAAAAAABAAEAAAICTAEAOw==';
                while (img.src == loading_gif) {
                    img.scrollIntoView({
                        behavior: 'smooth',
                        block: 'center',
                    });
                    await new Promise(resolve => setTimeout(resolve, 500));
                }
                images.push({ page_url, image_src: img.src, alt: img.alt });
            }
            return images;
        });
        return images;
    }
    async function scrollToBottom() {
        await page.evaluate(async () => {
            for (;;) {
                let imgs = document.querySelectorAll('[data-lpage] g-img img');
                let img = imgs[imgs.length - 1];
                if (!img)
                    return;
                img.scrollIntoView({ behavior: 'smooth', block: 'center' });
                await new Promise(resolve => setTimeout(resolve, 500));
                let bars = document.querySelectorAll('[role="progressbar"]');
                let bar = bars[bars.length - 1];
                if (!bar)
                    return;
                let rect = bar.getBoundingClientRect();
                let size = rect.width * rect.height;
                if (size == 0)
                    return;
                await new Promise(resolve => setTimeout(resolve, 500));
            }
        });
    }
    async function saveImage(image) {
        let { page_url, image_src, alt } = image;
        let res = await fetch(image_src);
        let mimeType = res.headers.get('Content-Type');
        if (!mimeType?.startsWith('image/')) {
            return;
        }
        let ext = mimeType.split('/')[1].split(';')[0];
        let binary = await res.arrayBuffer();
        let buffer = Buffer.from(binary);
        let hash = (0, crypto_1.createHash)('sha256');
        hash.write(buffer);
        let filename = hash.digest().toString('hex') + '.' + ext;
        let file = (0, path_1.join)(dir, filename);
        let row = (0, better_sqlite3_proxy_1.find)(proxy_1.proxy.image, { filename });
        if (row) {
            return;
        }
        let fileSize = await getFileSize(file);
        if (fileSize != buffer.length) {
            await (0, promises_1.writeFile)(file, buffer);
        }
        let domain = new URL(page_url).hostname;
        let domain_id = (0, better_sqlite3_proxy_1.seedRow)(proxy_1.proxy.domain, { domain });
        let page_id = (0, better_sqlite3_proxy_1.seedRow)(proxy_1.proxy.page, { url: page_url }, { domain_id });
        proxy_1.proxy.image.push({ filename, page_id, keyword_id, alt, embedding: null });
    }
    let lastCount = 0;
    let attempt = 0;
    for (;;) {
        let images = await collectImages();
        let count = images.length;
        if (count != lastCount) {
            attempt = 0;
        }
        if (count > 0 && count == lastCount) {
            attempt++;
            if (attempt > 3) {
                break;
            }
        }
        cli.update(`${cli_prefix}searching "${keyword}": ${count} images ...`);
        for (let image of images.slice(lastCount)) {
            await saveImage(image);
        }
        await scrollToBottom();
        lastCount = count;
    }
    cli.update(`${cli_prefix}searched "${keyword}": ${lastCount} images.`);
    cli.nextLine();
}
async function getFileSize(file) {
    try {
        return (await (0, promises_1.stat)(file)).size;
    }
    catch (error) {
        return 0;
    }
}
async function main(keywords) {
    let n = keywords.length;
    for (let i = 0; i < n; i++) {
        let cli_prefix = `[${i + 1}/${n}] `;
        let keyword = keywords[i];
        if ((0, better_sqlite3_proxy_1.find)(proxy_1.proxy.keyword, { keyword })?.complete_time) {
            console.log(`${cli_prefix}skip "${keyword}"`);
            continue;
        }
        await collectByKeyword(keyword, { cli_prefix });
        (0, better_sqlite3_proxy_1.find)(proxy_1.proxy.keyword, { keyword }).complete_time = Date.now();
    }
    await (0, browser_1.closeBrowser)();
}