利用AI批量为二次元图片打上标签并存入数据库中

在前一篇文章的最后，我提到DeepDanbooru识别普通的二次元图片效果也不错，于是我就有了把自己收藏的所有图片全打上标签的想法，方便以后检索，也方便做出api。

数据库采用mongodb，语言自然还是python。

改造源码

前一篇文章中我们看到识别结果是一行行打印出来的，这自然不符合我们的储存标准，所以首先需要修改源码，把它的输出变为字典型，也即：image_name:tag_list。

首先找到分析函数，在DeepDanbooru\deepdanbooru\commands\evaluate_project.py里，把evaluate_project.py复制一份并重命名为evaluate.py，然后修改30行之后的代码，替换为：

result = {}
for image_path in taget_image_paths:
    image_name = image_path.split('\\')[-1]
    print('正在分析{0}中……'.format(image_name))
    try:
        image = dd.data.load_image_for_evaluate(
            image_path, width=width, height=height)

        image_shape = image.shape
        image = image.reshape(
            (1, image_shape[0], image_shape[1], image_shape[2]))
        y = model.predict(image)[0]

        tag_list = []

        for i, tag in enumerate(tags):
            if y[i] < threshold:
                continue
            tag_list.append(tag)
      
        result[image_name] = tag_list
    except Exception:
        print('{0}无法识别'.format(image_name))
        result[image_name] = '000000'
return result

[scode type="yellow"]不要忘了把函数名改为evaluate[/scode]

然后按照https://www.sitstars.com/archives/73/所言，修改DeepDanbooru\deepdanbooru\commands\__init__.py，添加一行代码：

from .evaluate import evaluate

这样我们就DeepDanbooru模块中在添加了一个evaluate函数，主要功能就是分析文件夹中的图片并返回标签。

导入模块并进行分析、存储

思路是这样的：因为源码是分析整个文件夹，所以我们要输入一个含有图片的待检测文件夹（当然你可以直接在源码中写函数，然后一个个分析一个个存储，只是从复用的角度来讲我觉得还是尽量保持源码的纯洁性比较好）。因为我的图片比较多，所以我首先考虑的是效率，已经贴过标签的图片就不要再重复分析了。所以对于每张图，首先查询一下数据库，看是否有'tag'，没有的话复制到临时文件夹中。所有图片检测过之后，就分析临时文件夹中的图，并更新相应的数据。如果所有的图片均已有tag，那么就没有必要运行DeepDanbooru了

另外，因为该程序无法分析gif格式动图，同时为了防止有其他格式的文件混入导致程序运行失败，因此加上后缀名检测，除了'jpg','png','jpeg'格式外一律跳过。

哦对了，因为我的数据库中已经有这些图片的其他数据（图片名、hash值、链接等），所以用的是update_one函数，而且没有插入图片名，如果你有相关需求，自行修改为insert_one

import sys
sys.path.append(r"E:\deeplearn\DeepDanbooru")
from deepdanbooru import commands
import time
import pymongo
import os,shutil

myclient = pymongo.MongoClient("mongodb://localhost:27017/")  # 连接数据库
mydb = myclient["yourdb"]   # 创建数据库
mycol = mydb["yourcol"]  # 创建集合

def mult_tag(pic_dir,project_dir):  
    print('{0}正在执行贴标签程序{0}'.format('*'*4))
    print('{0}待识别图片整理中……{0}'.format('*'*4))
    condition = 0
    for pic in os.listdir(pic_dir):
        last = pic[pic.rindex(r'.')+1:] # 获取后缀
        if last not in ['jpg','png','jpeg']: # 不能识别gif等格式
            continue
        data = mycol.find_one({"name": pic})
        if data:
            try: 
                data['tag']
                print('{0}跳过'.format(pic))
            except Exception:
                condition = 1
                temp_dir = os.path.dirname(pic_dir)+'\\temp' # 父文件夹加上temp
  
                if os.path.exists(temp_dir) == False:
                    os.mkdir(temp_dir)
                pic_path = os.path.join(pic_dir, pic)  
                pic_path_temp = os.path.join(temp_dir, pic)  
                shutil.copy(pic_path,pic_path_temp)
                print('{0}复制到temp文件中'.format(pic))
  
        else:
            print('{0}未上传，请注意！'.format(pic))
            time.sleep(10)
  
    if condition == 0:
        print('所有图片均已贴上标签，无需运行程序')
    else:
        print('{0}智能识别中……{0}'.format('*'*4))
        tag_sum = commands.evaluate(project_dir, temp_dir, 0.5)
        print('智能识别完毕，存入数据库中……')
        for pic in os.listdir(temp_dir):
            mycol.update_one({"name":pic},{"$set":{"tag":tag_sum[pic]}})  
        print('全部存入数据库，共为{0}张图片贴上标签'.format(len(tag_sum)))
        print('准备删除temp目录下所有图片……')
        for pic in os.listdir(temp_dir):
            pic_path = os.path.join(temp_dir,pic)
            os.remove(pic_path) 
        print('temp文件夹下图片删除完毕')

if __name__ == '__main__':
    pic_dir = r'E:\壁纸\已上传壁纸'  # 要上传的文件夹
    project_dir = 'E:\\deeplearn\\DeepDanbooru\\project' # 预训练模型文件夹
    mult_tag(pic_dir,project_dir)

翻译（未完成）

DeepDanbooru给出的标签全是英文的（还有很多罗马音），对于英语渣的我自然无法忍受，因此就有了把它汉化的想法。自然不可能手动一个个翻，我的想法是先借助翻译平台进行翻译，然后再进行修改。之前注册过百度翻译开放平台，就选它了。

我们先读取tag文件（在DeepDanbooru\project\tags.txt中），返回一个列表，然后进行翻译并储存到数据库中。因为tag比较多，为了防止网络原因意外中断，也要加个数据库检测以免重复工作。

首先写百度翻译函数，官网提供了python demo，稍微修改一下就能用。

import http.client
import hashlib
import urllib
import random
import json

def bd(text,fromLang = 'auto',toLang = 'zh'):
    appid = ''  # 填写你的appid
    secretKey = ''  # 填写你的密钥  
    httpClient = None
    myurl = '/api/trans/vip/translate'

    salt = random.randint(32768, 65536)
    sign = appid + text + str(salt) + secretKey
    sign = hashlib.md5(sign.encode()).hexdigest()
    myurl = myurl + '?appid=' + appid + '&q=' + urllib.parse.quote(text) + '&from=' + fromLang + '&to=' + toLang + '&salt=' + str(
    salt) + '&sign=' + sign
  
    try:
        httpClient = http.client.HTTPConnection('api.fanyi.baidu.com')
        httpClient.request('GET', myurl)
  
        # response是HTTPResponse对象
        response = httpClient.getresponse()
        result_all = response.read().decode("utf-8")
        result = json.loads(result_all)
        return result

    except Exception as e:
        print (e)
    finally:
        if httpClient:
            httpClient.close()

然后是读取tag文件并批量翻译：

def transalate(txt_path):
    mycol = mydb["tags"] 
    print('{0}正在翻译中……{0}'.format('*'*4))
    with open(txt_path,'r') as f:
        tag_list = f.readlines()
        for tag in tag_list:
            tag = tag.replace('\n','') # 删掉换行符
            if mycol.find_one({'en': tag}):
                print('{0}已翻译，跳过'.format(tag))
            else:
                tag_org = tag
                tag = tag.replace('_',' ') # 为了翻译效果删掉_分隔符
                result = trans.bd(tag)
                cn = result['trans_result'][0]['dst']
                mycol.insert_one({'en':tag_org,'cn':cn}) 
                print('{0}翻译为：{1}'.format(tag,cn))
    print('{0}翻译完毕，结果已经保存至数据库'.format('*'*4))

def main():
    txt_path = r'E:\deeplearn\DeepDanbooru\project\tags.txt'
    transalate(txt_path)

为了方便修改，你可以先导出为json格式，然后用vscode打开，修改完毕后再导入。

mongoexport -d mdpicture -c tags -o D:\mogodata\export\tags.json # 导出
mongoimport -d mdpicture -c tags --file D:\mogodata\export\tags.json # 导入

效果如下图（未修改）：

因为量实在太大（7K多条），而且好多英文（罗马音？）我不太明白意思，百度翻译出来的效果也不太好，所以汉化计划只能慢慢做了。

版权属于：作者名称

本文链接：https://www.sitstars.com/archives/75/

转载时须注明出处及本声明