Spaces:
Running
Running
| import Koa from 'koa' | |
| import http from 'axios' | |
| import sharp from 'sharp' | |
| import tf from '@tensorflow/tfjs-node' | |
| import nsfw from 'nsfwjs' | |
| import { extractColors } from './extract-colors.js' | |
| import FastAverageColor from './fast-average-color.js' | |
| import { phash } from './phash.js' | |
| const app = new Koa() | |
| app.use(async (ctx) => { | |
| const { data: image_buffer } = await http({ | |
| method: 'GET', | |
| url: String(ctx.request.query.url).replace( | |
| 'i.pximg.net', | |
| 'i.pixiv.re', | |
| ), | |
| responseType: 'arraybuffer', | |
| }) | |
| const { data: image_raw, info } = await sharp(image_buffer, { | |
| animated: true, | |
| }) | |
| .ensureAlpha() | |
| .raw() | |
| .toBuffer({ resolveWithObject: true }) | |
| const image_phash = await phash(image_buffer) | |
| const colors = await extractColors( | |
| { data: image_raw, width: info.width, height: info.height }, | |
| { | |
| pixels: info.width * info.height, | |
| distance: 0.05, | |
| splitPower: 6, | |
| colorValidator: (red, green, blue, alpha = 255) => alpha >= 255, | |
| hueDistance: 0, | |
| saturationDistance: 1, | |
| lightnessDistance: 1, | |
| }, | |
| ) | |
| const fac = new FastAverageColor() | |
| const { | |
| value: [r, g, b, a], | |
| hex, | |
| isLight: blight, | |
| } = fac.prepareResult(fac.getColorFromArray4(image_raw)) | |
| const image_resize = await sharp(image_buffer, { animated: true }) | |
| .resize({ | |
| width: 128, | |
| height: 128, | |
| fit: 'inside', | |
| }) | |
| .gif({ dither: 0 }) | |
| .toBuffer() | |
| // const model = await nsfw.load(`https://raw.githubusercontent.com/infinitered/nsfwjs/refs/heads/master/models/mobilenet_v2_mid/model.json`, { type: 'graph' }) | |
| // // const model = await nsfw.load(`https://raw.githubusercontent.com/infinitered/nsfwjs/refs/heads/master/models/mobilenet_v2/model.json`) | |
| // // const model = await nsfw.load(`https://s3.amazonaws.com/ir_public/nsfwjscdn/TFJS_nsfw_mobilenet/tfjs_quant_nsfw_mobilenet/model.json`) | |
| const model = await nsfw.load(`https://raw.githubusercontent.com/infinitered/nsfwjs/refs/heads/master/models/inception_v3/model.json`, { size: 299 }) | |
| const predictions = await model.classifyGif(image_resize, { fps: 1 }) | |
| const detail = new Map() | |
| predictions | |
| .flat(Infinity) | |
| .sort((a, b) => a.probability - b.probability) | |
| .forEach((i) => { | |
| detail.set( | |
| i.className.toLocaleLowerCase(), | |
| Math.floor(i.probability * 100), | |
| ) | |
| }) | |
| const fail = ['hentai', 'porn', 'sexy'] | |
| const result = [...detail] | |
| .filter((i) => i.pop() > 7) | |
| .some((i) => fail.includes(i.pop())) | |
| tf.dispose() | |
| model.model.dispose() | |
| ctx.body = { | |
| nsfw: Object.assign({ result }, Object.fromEntries(detail)), | |
| colors: colors | |
| .sort((a, b) => b.area - a.area) | |
| .slice(0, 3) | |
| .map((i) => { | |
| // hex: `#${Number(十进制).toString(16).padStart(6, 0)}` | |
| return { | |
| color: parseInt(String(i.hex).replace('#', ''), 16), | |
| ratio: Number(i.area.toFixed(2)), | |
| } | |
| }), | |
| color: hex, | |
| // blight, | |
| phash: image_phash, | |
| } | |
| }) | |
| app.listen(process.env.PORT || 7860) |