Spaces:
Sleeping
Sleeping
rogerxavier
commited on
Commit
•
b0825a9
1
Parent(s):
8416394
Update utils.py
Browse files
utils.py
CHANGED
@@ -1,246 +1,261 @@
|
|
1 |
-
import json
|
2 |
-
import os
|
3 |
-
import random
|
4 |
-
from zipfile import ZipFile
|
5 |
-
import zipfile
|
6 |
-
import gradio as gr
|
7 |
-
from config import *
|
8 |
-
import shutil
|
9 |
-
from PIL import Image#上传保存遮罩和比较遮罩width height能否和章节图片匹配的时候用
|
10 |
-
from taskMap import *
|
11 |
-
|
12 |
-
|
13 |
-
def load_config():
|
14 |
-
return get_variables()
|
15 |
-
|
16 |
-
class configData():
|
17 |
-
def __init__(self):
|
18 |
-
self.data = load_config()
|
19 |
-
def update(self):
|
20 |
-
self.data = load_config()
|
21 |
-
print(self.data)
|
22 |
-
|
23 |
-
configData = configData()#复用实例
|
24 |
-
|
25 |
-
|
26 |
-
# 解压缩函数
|
27 |
-
def unzip_file(file_path, extract_path):
|
28 |
-
with zipfile.ZipFile(file_path, 'r') as zip_ref:
|
29 |
-
zip_ref.extractall(extract_path)
|
30 |
-
|
31 |
-
# 上传并解压缩函数
|
32 |
-
#upload_and_unzip(file:"文件",extract_path:"解压路径")->"void 成功提示":
|
33 |
-
def upload_and_unzip(file_obj, extract_path):
|
34 |
-
print("执行上传解压函数")
|
35 |
-
##这样写如果新的多,那么替换为多的一方,如果新的少,也只会替换已经存在的部分--->符合预期想法
|
36 |
-
with ZipFile(file_obj.name) as zfile:
|
37 |
-
zfile.extractall(extract_path)
|
38 |
-
return "File uploaded and extracted successfully to "+extract_path
|
39 |
-
|
40 |
-
#构造函数获取当前目录file list信息-不放config进行返回是为了避免引入旧值
|
41 |
-
def update_file_info(root_dir):
|
42 |
-
info = {}
|
43 |
-
for root, dirs, files in os.walk(root_dir):
|
44 |
-
info[root] = {
|
45 |
-
"directories": dirs,
|
46 |
-
"files": files
|
47 |
-
}
|
48 |
-
for dir_name in dirs:
|
49 |
-
update_file_info(os.path.join(root, dir_name))
|
50 |
-
return info
|
51 |
-
|
52 |
-
def random_chapter():
|
53 |
-
# for root, info in configData.data['file_info'].items():
|
54 |
-
#手动更新避免引入旧值-成功
|
55 |
-
file_info = update_file_info(configData.data['current_dir'])
|
56 |
-
configData.data['file_info'] = file_info
|
57 |
-
for root, info in configData.data['file_info'].items():
|
58 |
-
if root.startswith(configData.data['manga_abs_dir']):
|
59 |
-
for directory in info["directories"]:
|
60 |
-
if any(file.endswith('.jpg') or file.endswith('.png') for file in os.listdir(os.path.join(root, directory))):
|
61 |
-
chapter_path = os.path.join(root, directory)
|
62 |
-
# 删除章节目录
|
63 |
-
print("章节目录是:",chapter_path)
|
64 |
-
|
65 |
-
#获取标题
|
66 |
-
# 将路径转换为标题
|
67 |
-
folders = chapter_path.split(os.sep)
|
68 |
-
if len(folders) >= 2:
|
69 |
-
chapter_title = " ".join([folders[-2], folders[-1]]) # 获取倒数第二个和倒数第一个目录名作为标题
|
70 |
-
else:
|
71 |
-
chapter_title = chapter_path # 如果目录层级不足2层,直接使用路径作为标题
|
72 |
-
|
73 |
-
# 获取标题
|
74 |
-
return chapter_title,chapter_path#返回章节标题和章节目录给任务-1批量顺序上传原始图片到服务器manga.
|
75 |
-
|
76 |
-
return None,None
|
77 |
-
|
78 |
-
def list_files():
|
79 |
-
return configData.data['file_info']
|
80 |
-
|
81 |
-
|
82 |
-
def save_mask(mask_file, output_dir):
|
83 |
-
#上传遮罩保存到mask下面,名称随意
|
84 |
-
# 检查遮罩文件是否存在
|
85 |
-
# 构造保存文件的路径-保证linux和windows都能转义
|
86 |
-
save_path = os.path.join(output_dir, '0.jpg')
|
87 |
-
# 创建目录(如果目录不存在)
|
88 |
-
os.makedirs(output_dir, exist_ok=True)
|
89 |
-
|
90 |
-
# 使用with open方式保存文件
|
91 |
-
mask_file_img = mask_file #这个是Image.open后的对象
|
92 |
-
mask_file_img.save(save_path)
|
93 |
-
return "文件保存成功到:" + str(save_path)
|
94 |
-
|
95 |
-
|
96 |
-
def update_config_json(newConfig:str):
|
97 |
-
newConfig = json.loads(newConfig) #转换str为dict
|
98 |
-
with open(configData.data['current_config_json'], "w") as file:
|
99 |
-
file.write(json.dumps(newConfig, indent=4, ensure_ascii=False)) # 指定indent参数来保持JSON格式
|
100 |
-
#返回新的json文件内容
|
101 |
-
with open(configData.data['current_config_json'], "r") as f:
|
102 |
-
content =f.read() # 读取文件内容
|
103 |
-
return content
|
104 |
-
|
105 |
-
|
106 |
-
def is_asp_task_valid()->bool:
|
107 |
-
# 启动状态检测 ->
|
108 |
-
# 1config允许定时任务 2 mask目录下齐备 3 manga_all目录下有可用素材 4 mask的width height和素材匹配
|
109 |
-
# 5应当将config写入文件而不是py随时可以修改,以防ck账号被封等情况下不发送->同时删除不可用账号 -全部space中账号反馈不可用的时候停止定时任务
|
110 |
-
# 6应该避免任务扎堆进行,对ocr space造成负担
|
111 |
-
if configData.data['allow_scheduler'] == False:
|
112 |
-
print("当前配置不允许定时任务,停止定时任务")
|
113 |
-
return False
|
114 |
-
chapter_title ,cur_chapter= random_chapter()
|
115 |
-
if cur_chapter == None:
|
116 |
-
print("当前没有可用章节了,停止定时任务")
|
117 |
-
return False
|
118 |
-
mask_file_path = os.path.join(configData.data['mask_dir'], '0.jpg')
|
119 |
-
if not os.path.exists(mask_file_path):
|
120 |
-
print("遮罩文件不存在,停止定时任务")
|
121 |
-
return False
|
122 |
-
# 读取遮罩图片长宽比较当前章节下面图片的长宽
|
123 |
-
mask_image = Image.open(mask_file_path)
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
if
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
189 |
-
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
229 |
-
|
230 |
-
|
231 |
-
|
232 |
-
|
233 |
-
|
234 |
-
|
235 |
-
|
236 |
-
|
237 |
-
|
238 |
-
|
239 |
-
|
240 |
-
|
241 |
-
|
242 |
-
|
243 |
-
|
244 |
-
|
245 |
-
|
246 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
import os
|
3 |
+
import random
|
4 |
+
from zipfile import ZipFile
|
5 |
+
import zipfile
|
6 |
+
import gradio as gr
|
7 |
+
from config import *
|
8 |
+
import shutil
|
9 |
+
from PIL import Image#上传保存遮罩和比较遮罩width height能否和章节图片匹配的时候用
|
10 |
+
from taskMap import *
|
11 |
+
|
12 |
+
|
13 |
+
def load_config():
|
14 |
+
return get_variables()
|
15 |
+
|
16 |
+
class configData():
|
17 |
+
def __init__(self):
|
18 |
+
self.data = load_config()
|
19 |
+
def update(self):
|
20 |
+
self.data = load_config()
|
21 |
+
print(self.data)
|
22 |
+
|
23 |
+
configData = configData()#复用实例
|
24 |
+
|
25 |
+
|
26 |
+
# 解压缩函数
|
27 |
+
def unzip_file(file_path, extract_path):
|
28 |
+
with zipfile.ZipFile(file_path, 'r') as zip_ref:
|
29 |
+
zip_ref.extractall(extract_path)
|
30 |
+
|
31 |
+
# 上传并解压缩函数
|
32 |
+
#upload_and_unzip(file:"文件",extract_path:"解压路径")->"void 成功提示":
|
33 |
+
def upload_and_unzip(file_obj, extract_path):
|
34 |
+
print("执行上传解压函数")
|
35 |
+
##这样写如果新的多,那么替换为多的一方,如果新的少,也只会替换已经存在的部分--->符合预期想法
|
36 |
+
with ZipFile(file_obj.name) as zfile:
|
37 |
+
zfile.extractall(extract_path)
|
38 |
+
return "File uploaded and extracted successfully to "+extract_path
|
39 |
+
|
40 |
+
#构造函数获取当前目录file list信息-不放config进行返回是为了避免引入旧值
|
41 |
+
def update_file_info(root_dir):
|
42 |
+
info = {}
|
43 |
+
for root, dirs, files in os.walk(root_dir):
|
44 |
+
info[root] = {
|
45 |
+
"directories": dirs,
|
46 |
+
"files": files
|
47 |
+
}
|
48 |
+
for dir_name in dirs:
|
49 |
+
update_file_info(os.path.join(root, dir_name))
|
50 |
+
return info
|
51 |
+
|
52 |
+
def random_chapter():
|
53 |
+
# for root, info in configData.data['file_info'].items():
|
54 |
+
#手动更新避免引入旧值-成功
|
55 |
+
file_info = update_file_info(configData.data['current_dir'])
|
56 |
+
configData.data['file_info'] = file_info
|
57 |
+
for root, info in configData.data['file_info'].items():
|
58 |
+
if root.startswith(configData.data['manga_abs_dir']):
|
59 |
+
for directory in info["directories"]:
|
60 |
+
if any(file.endswith('.jpg') or file.endswith('.png') for file in os.listdir(os.path.join(root, directory))):
|
61 |
+
chapter_path = os.path.join(root, directory)
|
62 |
+
# 删除章节目录
|
63 |
+
print("章节目录是:",chapter_path)
|
64 |
+
|
65 |
+
#获取标题
|
66 |
+
# 将路径转换为标题
|
67 |
+
folders = chapter_path.split(os.sep)
|
68 |
+
if len(folders) >= 2:
|
69 |
+
chapter_title = " ".join([folders[-2], folders[-1]]) # 获取倒数第二个和倒数第一个目录名作为标题
|
70 |
+
else:
|
71 |
+
chapter_title = chapter_path # 如果目录层级不足2层,直接使用路径作为标题
|
72 |
+
|
73 |
+
# 获取标题
|
74 |
+
return chapter_title,chapter_path#返回章节标题和章节目录给任务-1批量顺序上传原始图片到服务器manga.
|
75 |
+
|
76 |
+
return None,None
|
77 |
+
|
78 |
+
def list_files():
|
79 |
+
return configData.data['file_info']
|
80 |
+
|
81 |
+
|
82 |
+
def save_mask(mask_file, output_dir):
|
83 |
+
#上传遮罩保存到mask下面,名称随意
|
84 |
+
# 检查遮罩文件是否存在
|
85 |
+
# 构造保存文件的路径-保证linux和windows都能转义
|
86 |
+
save_path = os.path.join(output_dir, '0.jpg')
|
87 |
+
# 创建目录(如果目录不存在)
|
88 |
+
os.makedirs(output_dir, exist_ok=True)
|
89 |
+
|
90 |
+
# 使用with open方式保存文件
|
91 |
+
mask_file_img = mask_file #这个是Image.open后的对象
|
92 |
+
mask_file_img.save(save_path)
|
93 |
+
return "文件保存成功到:" + str(save_path)
|
94 |
+
|
95 |
+
|
96 |
+
def update_config_json(newConfig:str):
|
97 |
+
newConfig = json.loads(newConfig) #转换str为dict
|
98 |
+
with open(configData.data['current_config_json'], "w") as file:
|
99 |
+
file.write(json.dumps(newConfig, indent=4, ensure_ascii=False)) # 指定indent参数来保持JSON格式
|
100 |
+
#返回新的json文件内容
|
101 |
+
with open(configData.data['current_config_json'], "r") as f:
|
102 |
+
content =f.read() # 读取文件内容
|
103 |
+
return content
|
104 |
+
|
105 |
+
|
106 |
+
def is_asp_task_valid()->bool:
|
107 |
+
# 启动状态检测 ->
|
108 |
+
# 1config允许定时任务 2 mask目录下齐备 3 manga_all目录下有可用素材 4 mask的width height和素材匹配
|
109 |
+
# 5应当将config写入文件而不是py随时可以修改,以防ck账号被封等情况下不发送->同时删除不可用账号 -全部space中账号反馈不可用的时候停止定时任务
|
110 |
+
# 6应该避免任务扎堆进行,对ocr space造成负担
|
111 |
+
if configData.data['allow_scheduler'] == False:
|
112 |
+
print("当前配置不允许定时任务,停止定时任务")
|
113 |
+
return False
|
114 |
+
chapter_title ,cur_chapter= random_chapter()
|
115 |
+
if cur_chapter == None:
|
116 |
+
print("当前没有可用章节了,停止定时任务")
|
117 |
+
return False
|
118 |
+
mask_file_path = os.path.join(configData.data['mask_dir'], '0.jpg')
|
119 |
+
if not os.path.exists(mask_file_path):
|
120 |
+
print("遮罩文件不存在,停止定时任务")
|
121 |
+
return False
|
122 |
+
# 读取遮罩图片长宽比较当前章节下面图片的长宽
|
123 |
+
mask_image = Image.open(mask_file_path)
|
124 |
+
# 获取指定目录下所有文件
|
125 |
+
all_files = os.listdir(cur_chapter)
|
126 |
+
# 过滤出图片文件
|
127 |
+
image_files = [file for file in all_files if file.endswith(('.jpg', '.jpeg', '.png', '.gif'))]
|
128 |
+
if image_files:
|
129 |
+
# 随机选择一个图片文件
|
130 |
+
random_chapter_image_name = random.choice(image_files)
|
131 |
+
random_chapter_image_path = os.path.join(cur_chapter, random_chapter_image_name)
|
132 |
+
|
133 |
+
# 打开选定的图片文件
|
134 |
+
random_chapter_image = Image.open(random_chapter_image_path)
|
135 |
+
|
136 |
+
# 现在random_chapter_image中包含了选中的随机图片文件
|
137 |
+
if mask_image.size != random_chapter_image.size:
|
138 |
+
print("遮罩大小和章节随机一个图片的大小不匹配,停止定时任务")
|
139 |
+
return False
|
140 |
+
else:
|
141 |
+
print("No image files found in the directory.")
|
142 |
+
return False
|
143 |
+
|
144 |
+
|
145 |
+
return True
|
146 |
+
|
147 |
+
#定时任务流程函数集合 ->返回None说明非正常执行情况
|
148 |
+
def run_asp_task():
|
149 |
+
if not is_asp_task_valid():
|
150 |
+
return None
|
151 |
+
#遍历bili_spaces中的space url并执行日常任务
|
152 |
+
for space_url in configData.data["bili_spaces"]:
|
153 |
+
process_task_list("0 1 2",space_url)
|
154 |
+
|
155 |
+
|
156 |
+
|
157 |
+
|
158 |
+
|
159 |
+
|
160 |
+
|
161 |
+
#实现gradio接受上传压缩文件,解压到manga_all目录下面
|
162 |
+
with gr.Blocks() as mangaManager:
|
163 |
+
#上传文件选择列表
|
164 |
+
file_selected = gr.File(label="待上传压缩章节",file_types=['.zip', '.tar', '.gz'])
|
165 |
+
#上传zip的解压保存路径或mask的地址(必选) 根据需要选择
|
166 |
+
output_dir_gr = gr.Dropdown(label="上传zip的解压保存路径或mask的地址(必选)", choices=configData.data['default_values'])
|
167 |
+
#压缩包上传按钮
|
168 |
+
file_upload_btn = gr.Button("开始上传")
|
169 |
+
#获取到返回的结果-可以是上传的,可以是查看files信息,也可以是别的
|
170 |
+
someResult = gr.Textbox(label="获取按钮返回信息", type="text")
|
171 |
+
# 设置按钮点击事件(调用上传解压函数,将压缩包内容解压到指定目录 默认是manga_all下面)
|
172 |
+
file_upload_btn.click(fn=upload_and_unzip, inputs=[file_selected, output_dir_gr],outputs=someResult)
|
173 |
+
|
174 |
+
#设置按钮查看json类型的listFiles信息
|
175 |
+
|
176 |
+
filesInfoBtn = gr.Button("查看files信息")
|
177 |
+
filesInfoBtn.click(fn=list_files,outputs=someResult)
|
178 |
+
|
179 |
+
#设置mask上传遮罩按钮保存遮罩到mask目录-因为直接重启space利用docker上传会触发定时任务,也许造成不好结果
|
180 |
+
# mask_selected = gr.inputs.Image(label="待上传遮罩", type='pil')
|
181 |
+
mask_selected = gr.components.Image(label="待上传遮罩", type='pil')#代替inputs的新写法
|
182 |
+
markUploadBtn = gr.Button("上传遮罩mask")
|
183 |
+
markUploadBtn.click(fn=save_mask,inputs=[mask_selected, output_dir_gr],outputs=someResult)
|
184 |
+
|
185 |
+
#设置按钮关联新config输入进行更新操作
|
186 |
+
config_update_text = gr.Textbox(placeholder="输入新的JSON数据")#str类型
|
187 |
+
config_update_btn = gr.Button("更新config.json数据")
|
188 |
+
config_update_btn.click(fn=update_config_json, inputs=[config_update_text], outputs=someResult)
|
189 |
+
|
190 |
+
|
191 |
+
# 定义一个函数来处理输入的步骤转list然后执行,返回执行结果->taskResult/None
|
192 |
+
# gradio输入0 1 2 3就行,会自动转"0 1 2 3"然后给process_task_list处理
|
193 |
+
def process_task_list(input_task_str:str = None,baseUrl:str =None):
|
194 |
+
if not is_asp_task_valid():
|
195 |
+
print("任务基本条件不符合")
|
196 |
+
return None
|
197 |
+
if baseUrl is None:
|
198 |
+
print("baseUrl不存在")
|
199 |
+
return None
|
200 |
+
if input_task_str is None:
|
201 |
+
print("input_task_str不存在")
|
202 |
+
return None
|
203 |
+
chapter_title,cur_chapter_path = random_chapter()
|
204 |
+
task_results = {}#保存任务执行结果
|
205 |
+
#临时变量记录是���上传了章节-上传了需要删除,而且只能在for循环任务结束后删除
|
206 |
+
manga_has_uploaded = False
|
207 |
+
|
208 |
+
|
209 |
+
print(input_task_str)# "a b c" - >['a', 'b', 'c']
|
210 |
+
task_list = input_task_str.split(' ')
|
211 |
+
for task in task_list:
|
212 |
+
if task in task_functions:
|
213 |
+
task_func = task_functions[task]#func
|
214 |
+
task_param = tasks_params[task]#dict
|
215 |
+
if "baseUrl" in task_param:
|
216 |
+
task_param["baseUrl"] = baseUrl
|
217 |
+
if "mask_path" in task_param:
|
218 |
+
#使用后遮罩不删,否则维护成本太高
|
219 |
+
task_param["mask_path"] = configData.data['mask_abs_dir']
|
220 |
+
if "manga_path" in task_param:
|
221 |
+
task_param["manga_path"] = cur_chapter_path#上传一个章节作为本次素材,保存到space的manga下面
|
222 |
+
print("上传了章节:",cur_chapter_path,"下面提交后删除")
|
223 |
+
manga_has_uploaded = True
|
224 |
+
#执行了这个任务就要删除该章节
|
225 |
+
|
226 |
+
if "bili_meta" in task_param:
|
227 |
+
bili_meta_data["title"] = chapter_title
|
228 |
+
task_param["bili_meta"] = bili_meta_data
|
229 |
+
result = task_func(**task_param)#执行对应函数获取结果
|
230 |
+
task_results[task] = "success" if result is None else result#保存对应函数执行结果
|
231 |
+
else:
|
232 |
+
print("该执行流程函数不存在")
|
233 |
+
return None#立即推出循环结束函数
|
234 |
+
if manga_has_uploaded is True:
|
235 |
+
#如果上传章节任务得到了执行,那么删除
|
236 |
+
shutil.rmtree(cur_chapter_path) # 递归地删除指定路径的目录及其所有内容,包括文件和子目录 #删除应该在返回的时候删
|
237 |
+
|
238 |
+
|
239 |
+
|
240 |
+
return task_results
|
241 |
+
|
242 |
+
# 创建一个输入块,接受str+空格类型的输入,比如设置默认值3只查看结果output.mp4
|
243 |
+
|
244 |
+
# 将按钮添加到任务管理器中
|
245 |
+
with gr.Blocks() as taskManager:
|
246 |
+
# 任务需要 1:查看指定space的output video状态 2: 手动修改执行流程->比如上传后是发送还是查看output效果
|
247 |
+
# 3:可以添加按钮跳转路由
|
248 |
+
task_list_text = gr.components.Textbox(type="text", label="输入任务列表元素,用空格分隔(一般用来执首尾比如查看output状态)", lines=5)
|
249 |
+
baseSpaceUrl = gr.components.Textbox(type="text", label="指定手动任务的执行容器地址", lines=3)
|
250 |
+
# 获取到返回的结果-可以是上传的,可以是查看files信息,也可以是别的
|
251 |
+
someResult = gr.components.Textbox(label="获取按钮返回信息", type="text")
|
252 |
+
# 创建一个按钮块来触发处理函数
|
253 |
+
taskBtn = gr.Button("处理列表")
|
254 |
+
taskBtn.click(fn=process_task_list,inputs=[task_list_text, baseSpaceUrl],outputs=someResult)
|
255 |
+
new_text = gr.components.Textbox(placeholder="敬请期待") # str类型
|
256 |
+
|
257 |
+
|
258 |
+
|
259 |
+
|
260 |
+
|
261 |
+
|