feat: support AnimeTrace (#34)

package.json

@@ -45,11 +45,12 @@
     "koishi": "^4.17.4"
   },
   "dependencies": {
-    "cheerio": "^1.0.0-rc.12",
+    "cheerio": "^1.0.0",
     "iqdb-client": "^3.0.0",
     "nhentai-api": "^3.4.3"
   },
   "devDependencies": {
+    "@koishijs/canvas": "^0.2.0",
     "@types/node": "^20.12.7",
     "koishi": "^4.17.4",
     "typescript": "^5.4.5"

src/animetrace.ts

@@ -0,0 +1,103 @@
+import { Context, h, HTTP, Session } from 'koishi'
+import type { Image } from '@koishijs/canvas'
+
+interface AnimeTraceRequest {
+  model: AnimeTraceModel
+  ai_detect: 1 | 0
+  force_one: 1 | 0
+}
+
+type AnimeTraceModel =
+  | 'large_model_preview'
+  | 'anime'
+  | 'anime_model_lovelive'
+  | 'pre_stable'
+  | 'game'
+  | 'game_model_kirakira'
+
+interface AnimeTraceResponse {
+  ai: boolean
+  code: number
+  data: ResponseData[]
+  new_code: number
+  msg?: string
+}
+
+interface ResponseData {
+  box: number[]
+  char: ResponseDataChar[]
+  box_id: string
+}
+
+interface ResponseDataChar {
+  name: string
+  cartoonname: string
+  acc: number
+}
+
+async function crop(ctx: Context, image: Image, box: ResponseData['box']): Promise<string> {
+  const width: number = image.naturalWidth ?? image['width']
+  const height: number = image.naturalHeight ?? image['height']
+  const outputWidth = width * (box[2] - box[0])
+  const outputHeight = height * (box[3] - box[1])
+  const canvas = await ctx.canvas.createCanvas(outputWidth, outputHeight)
+  canvas.getContext('2d').drawImage(
+    image,
+    width * box[0],
+    height * box[1],
+    outputWidth,
+    outputHeight,
+    0,
+    0,
+    outputWidth,
+    outputHeight
+  )
+  return await canvas.toDataURL('image/png')
+}
+
+async function makeSearch(http: HTTP, url: string, ctx: Context): Promise<string> {
+  const { data, filename, type } = await http.file(url)
+  const form = new FormData()
+  const value = new Blob([data], { type })
+  form.append('image', value, filename)
+
+  const res = await http.post<AnimeTraceResponse>('https://aiapiv2.animedb.cn/ai/api/detect', form, {
+    responseType: 'json',
+    params: {
+      force_one: 1,
+      model: 'anime_model_lovelive',
+      ai_detect: 0
+    } as AnimeTraceRequest
+  })
+
+  if (res.code !== 0 && res.msg) {
+    return '搜图时遇到问题：' + res.msg
+  } else if (res.data.length === 0) {
+    return '没有识别到角色'
+  }
+
+  const image = await ctx.canvas.loadImage(data)
+  const elements = []
+  for (const v of res.data) {
+    elements.push(
+      h.image(await crop(ctx, image, v.box)),
+      h.text(`角色：${v.char[0].name}`),
+      h.text(`来源：${v.char[0].cartoonname}`)
+    )
+  }
+  return elements.join('\n')
+}
+
+export default async function (http: HTTP, url: string, session: Session) {
+  let result = 'AnimeTrace 搜图\n'
+  try {
+    result += await makeSearch(http, url, session.app)
+  } catch (err) {
+    if (http.isError(err) && err.response.data?.msg) {
+      result += '搜图时遇到问题：' + err.response.data.msg
+    } else {
+      result += '搜图时遇到问题：' + err
+    }
+  }
+  return result
+}

src/ascii2d.ts

@@ -35,7 +35,7 @@ export default async function (http: Quester, url: string, session: Session, con
 }
 
 function getDetail(html: string, config: OutputConfig) {
-  const $ = load(html, { decodeEntities: false })
+  const $ = load(html)
   const $box = $($('.item-box')[1])
   if ($box.length === 0) {
     logger.warn('[error] ascii2d bovw cannot find images in web page')
@@ -56,6 +56,6 @@ function getDetail(html: string, config: OutputConfig) {
 }
 
 function getTokuchouUrl(html: string) {
-  const $ = load(html, { decodeEntities: false })
+  const $ = load(html)
   return `${baseURL}/search/bovw/${$($('.hash')[0]).text().trim()}`
 }

src/index.ts

@@ -1,13 +1,14 @@
-import { Awaitable, Command, Context, h, makeArray, Quester, Session } from 'koishi'
+import { Awaitable, Command, Context, Dict, h, makeArray, Quester, Session } from 'koishi'
 import ascii2d from './ascii2d'
 import saucenao from './saucenao'
 import iqdb from './iqdb'
+import animetrace from './animetrace'
 import { Config } from './utils'
 
 export { Config }
 
 export const name = 'image-search'
-export const inject = ['http']
+export const inject = ['http', 'canvas']
 
 async function mixedSearch(http: Quester, url: string, session: Session, config: Config) {
   return await saucenao(http, url, session, config, true) && ascii2d(http, url, session, config)
@@ -34,15 +35,17 @@ export function apply(ctx: Context, config: Config = {}) {
     return '令牌失效导致访问失败，请联系机器人作者。'
   })
 
-  ctx.command('search [image:text]', '搜图片')
+  ctx.command('search [image:image]', '搜图片')
     .shortcut('搜图', { fuzzy: true })
     .action(search(mixedSearch))
-  ctx.command('search/saucenao [image:text]', '使用 saucenao 搜图')
+  ctx.command('search/saucenao [image:image]', '使用 saucenao 搜图')
     .action(search(saucenao))
-  ctx.command('search/ascii2d [image:text]', '使用 ascii2d 搜图')
+  ctx.command('search/ascii2d [image:image]', '使用 ascii2d 搜图')
     .action(search(ascii2d))
-  ctx.command('search/iqdb [image:text]', '使用 iqdb 搜图')
+  ctx.command('search/iqdb [image:image]', '使用 iqdb 搜图')
     .action(search(iqdb))
+  ctx.command('search/animetrace [image:image]', '使用 animetrace 搜图')
+    .action(search(animetrace))
 
   const pendings = new Set<string>()
 
@@ -77,14 +80,13 @@ export function apply(ctx: Context, config: Config = {}) {
   }
 
   function search(callback: SearchCallback): Command.Action {
-    return async ({ session }, image) => {
+    return async ({ session }, image: Dict) => {
       const id = session.channelId
       if (pendings.has(id)) return '存在正在进行的查询，请稍后再试。'
 
-      const url = getUrl(image)
-      if (url) {
+      if (image?.src) {
         pendings.add(id)
-        return searchUrl(session, url, callback)
+        return searchUrl(session, image.src, callback)
       }
 
       const dispose = session.middleware(({ content }, next) => {