HTSD / GenerateTxT.配合danbooruDownloader图片抓取教程使用.py
xiuxiudesu's picture
add all_prompt generate
302e7bb
#!/usr/bin/python
print("me")
import os
import xml
import xml.sax
import sys
img_folder = "F:\StableDifussion\stable diffusion 工具\DanbooruDownloader3.20220828\images\gelbooru.com\deedlit\proccessed"
xml_path = "F:\StableDifussion\stable diffusion 工具\DanbooruDownloader3.20220828\savelist\deedlit.xml"
xml_imgs_list = [];
class MovieHandler( xml.sax.ContentHandler ):
def __init__(self):
self.CurrentData = ""
self.Id = ""
self.index = 0
self.Filename = ""
self.Tags = ""
self.row = []
# 元素开始事件处理
def startElement(self, tag, attributes):
self.CurrentData = tag
#print('self.CurrentData:'+str(self.CurrentData))
# 元素结束事件处理
def endElement(self, tag):
if self.CurrentData == "JpegHeight":
#xml_imgs_list.append(self.row)
#self.row = []
#self.index = self.index + 1
print ("JpegHeight:")
#print (self.type)
self.CurrentData = ""
# 内容事件处理
def characters(self, content):
if self.CurrentData == "Id":
print('Id:'+str(content))
#print('self.Id:'+str(content))
#self.Id =content
#print('xml_imgs_list:')
#print(xml_imgs_list)
self.row = []
self.row.append(content)
elif self.CurrentData == "Tags":
print('Tags:'+str(content))
#self.Tags = content
self.row.append(content)
elif self.CurrentData == "Filename":
print('Filename:'+str(content))
#print('self.Id:'+str(self.Id)+'\n')
#print('self.index:'+str(self.index)+'\n')
#self.Filename = content
self.row.append(content)
if len(self.row) < 3:
print(self.row)
print("该图片生成失败,按回车键继续")
input()
self.row = []
else:
xml_imgs_list.append(self.row)
self.row = []
parser = xml.sax.make_parser()
parser.setFeature(xml.sax.handler.feature_namespaces, 0)
Handler = MovieHandler()
parser.setContentHandler( Handler )
parser.parse(xml_path);
prompt_arr = []
for file_name in os.listdir(img_folder):
for xml_row in xml_imgs_list:
try:
if file_name in xml_row[2]:
text_content = xml_row[1]
prompt_arr.append(text_content)
file_path_explode_arr = os.path.splitext(file_name);
#save_path = img_folder+file_path_explode_arr[0]+'.txt'
save_path = os.path.join(img_folder,file_path_explode_arr[0]+".txt")
##print("save_path:"+save_path)
##print("\n\r")
##print("text_content:"+text_content)
##print("\n\r")
#fwrite(save_path,text_content);
file = open(save_path,'w+')
file.write(text_content);
file.close()
except:
print(xml_row)
print("该图片生成失败,按回车键继续")
input()
#print("OS error: {0}".format(err))
file = open(os.path.join(img_folder,"all_prompt.txt"),'w+')
file.write(",".join(prompt_arr));
file.close()