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

在[前一篇文章](https://www.sitstars.com/archives/72/)的最后，我提到`DeepDanbooru`识别普通的二次元图片效果也不错，于是我就有了把自己收藏的所有图片全打上标签的想法，方便以后检索，也方便做出api。

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

## 改造源码

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

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

“`python
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/](https://www.sitstars.com/archives/73/)所言，修改`DeepDanboorudeepdanboorucommands__init__.py`，添加一行代码：

“`python
from .evaluate import evaluate
“`

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

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

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

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

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

“`python
import sys
sys.path.append(r”E:deeplearnDeepDanbooru”)
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`给出的标签全是英文的（还有很多罗马音），对于英语渣的我自然无法忍受，因此就有了把它汉化的想法。自然不可能手动一个个翻，我的想法是先借助翻译平台进行翻译，然后再进行修改。之前注册过[百度翻译开放平台](https://api.fanyi.baidu.com)，就选它了。

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

首先写百度翻译函数，官网提供了python [demo](https://fanyiapp.bj.bcebos.com/api/demo/BaiduTransAPI-forPython3.py.zip)，稍微修改一下就能用。

“`python
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文件并批量翻译：

“`python
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:deeplearnDeepDanbooruprojecttags.txt’
transalate(txt_path)
“`

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

“`cmd
mongoexport -d mdpicture -c tags -o D:mogodataexporttags.json # 导出
mongoimport -d mdpicture -c tags –file D:mogodataexporttags.json # 导入
“`

效果如下图（未修改）：


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