@ibaraki-douji/waifu2x/src/Waifu2X.ts

Upscale image with waifu2x-ncnn-vulkan.

import { Options } from "./types";
import path = require("path");
import { tmpdir } from "os";
import { readFileSync, writeFileSync } from "fs";
import { ChildProcessWithoutNullStreams, spawn } from "child_process";

export = class Waifu2X {

    public readonly inputPath: string;
    public readonly outputPath: string;
    public readonly process: ChildProcessWithoutNullStreams;
    public endBuffer: Buffer;

    constructor (inp: string | Buffer, out: string, options: Options = {}) {
        if (typeof inp == "string") this.inputPath = path.resolve(inp);
        else {
            this.inputPath = path.resolve(tmpdir(), generate(50) + ".png");
            writeFileSync(this.inputPath, inp as Buffer);
        }
        if (options.outputAsBuffer) {
            this.outputPath = path.resolve(tmpdir(), generate(50) + ".png");
        } else {
            this.outputPath = path.resolve(out);
        }

        this.process = spawn(__dirname + "/../waifu2x/waifu2x-ncnn-vulkan", [
            "-i", this.inputPath,
            "-o", this.outputPath,
            ...(options.gpu ? ["-g", ""+options.gpu] : []),
            ...(options.noise ? ['-n', '' + options.noise] : []),
            ...(options.ramLimit ? ['-t', ''+options.ramLimit] : []),
            ...(options.scale ? ['-s', ''+options.scale] : [])
        ]);

        this.process.on('close', () => {
            try {
                if (options.outputAsBuffer) this.endBuffer = readFileSync(this.outputPath);
            } catch (e) {}
        })
    }

    public static upscale(inp: string | Buffer, out: string, options?: Options) {
        return new Waifu2X(inp, out, options);
    }

    public static async listGPUs() {
        const process = spawn(__dirname + "/../waifu2x/waifu2x-ncnn-vulkan", [
            "-i", "/" + generate(50) + ".png",
            "-o", "/" + generate(50) + ".png",
            "-s", "1",
        ]);
        const out = [];
        process.on('error', console.log)
        process.stdout.on('data', d => out.push(d.toString()));
        process.stderr.on('data', d => out.push(d.toString()));
        await new Promise<void>(res => (process.on('close', res)));
        const GPUs = [];
        for (let std of out) {
            if (std.startsWith("[")) {
                const str = std.split(']')[0].split("[")[1];
                GPUs[+str.split(" ")[0]] = str.replace(str.split(" ")[0] + " ", "");
            }
        }
        return GPUs
    }

    public finishedPromise() {
        return new Promise<void>(res => {
            process.on('exit', res)
        });
    }
}

const generate=(length)=>{const alphabet="AZERTYUIOPQSDFGHJKLMWXCVBNazertyuiopqsdfghjklmwxcvbn123456789";let end="";for(let i=0;i<length;i++)end+=alphabet[Math.floor(Math.random()*alphabet.length)];return end}