demo_gradio_batch.py 114 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786
  1. import os
  2. import io
  3. import uuid
  4. import json
  5. import zipfile
  6. import tempfile
  7. import threading
  8. import queue
  9. import shutil
  10. from pathlib import Path
  11. from PIL import Image
  12. import requests
  13. import gradio as gr
  14. import re
  15. import math
  16. import datetime
  17. # Local project imports (assumed available)
  18. from dots_ocr.utils import dict_promptmode_to_prompt
  19. from dots_ocr.utils.consts import MIN_PIXELS, MAX_PIXELS
  20. from dots_ocr.utils.demo_utils.display import read_image
  21. from dots_ocr.parser import DotsOCRParser
  22. # ---------------- Config & globals ----------------
  23. DEFAULT_CONFIG = {
  24. "ip": "127.0.0.1",
  25. "port_vllm": 8000,
  26. "min_pixels": MIN_PIXELS,
  27. "max_pixels": MAX_PIXELS,
  28. }
  29. # Absolute constraints discovered from runtime:
  30. ABS_MIN_PIXELS = 3136
  31. ABS_MAX_PIXELS = 11289600
  32. current_config = DEFAULT_CONFIG.copy()
  33. # default parser instance (can be overridden per-task)
  34. dots_parser = DotsOCRParser(
  35. ip=DEFAULT_CONFIG["ip"],
  36. port=DEFAULT_CONFIG["port_vllm"],
  37. dpi=200,
  38. min_pixels=DEFAULT_CONFIG["min_pixels"],
  39. max_pixels=DEFAULT_CONFIG["max_pixels"],
  40. )
  41. RESULTS_CACHE = {} # rid -> result dict or placeholder
  42. TASK_QUEUE = queue.Queue()
  43. # Worker pool for background processing (adjustable via UI)
  44. WORKER_THREADS = []
  45. MAX_CONCURRENCY = 6
  46. THREAD_LOCK = threading.Lock()
  47. RETRY_COUNTS = {} # rid -> attempts
  48. MAX_AUTO_RETRIES = 5
  49. RETRY_BACKOFF_BASE = 1.7
  50. DEFAULT_SCRIPT_TEMPLATE = """# 高级脚本使用说明
  51. # 提供对象: api
  52. # 日志: 使用 print(...) 或 debug(...) 输出到下方“脚本日志”实时区域。
  53. # api.get_ids() -> [rid,...] 按当前 UI 顺序返回
  54. # api.get_status(rid) -> {'status','ui': {'tab','nohf','source'}, 'filtered': bool, 'input_width': int, 'input_height': int}
  55. # api.get_texts(rid) -> {
  56. # 'md': 原始 Markdown, 'md_nohf': 原始 NOHF Markdown, 'json': 原始 JSON,
  57. # 'md_edit': 编辑版 Markdown 或 None, 'md_nohf_edit': 编辑版 NOHF Markdown 或 None, 'json_edit': 编辑版 JSON 或 None
  58. # }
  59. # api.choose_texts(rid, prefer_ui=True, prefer_edit=True, prefer_nohf=None) -> {'md','json'}
  60. # - prefer_ui: True 时根据当前 UI 的 NOHF/来源选择内容
  61. # - prefer_edit: True 时优先用编辑内容(若存在)
  62. # - prefer_nohf: 显式指定是否使用 NOHF(覆盖 UI),None 表示跟随 UI
  63. # api.list_paths(rid) -> {
  64. # 'temp_dir': str, 'session_id': str,
  65. # 'result': {'md':path,'md_nohf':path,'json':path,'layout':path or None,'image':path or None},
  66. # 'edited': {'md':path or None,'md_nohf':path or None,'json':path or None}
  67. # }
  68. # api.path_exists(path) -> bool 判断路径是否存在
  69. # api.build_export(name='custom') -> ExportBuilder
  70. # ExportBuilder:
  71. # .add_text('dir/file.md', '...') 写入文本
  72. # .add_bytes('bin/data.bin', b'...') 写入二进制
  73. # .add_file('/abs/path/file.md', 'dir/file.md') 拷贝已有文件
  74. # .mkdir('subdir/') 创建目录
  75. # .finalize() -> zip_path 打包为 zip 并返回路径
  76. #
  77. # 约定: 定义 main(api) 并返回以下之一:
  78. # - ExportBuilder 实例(将自动 finalize)
  79. # - 目录路径或文件路径(目录将被打包为 zip)
  80. # - None(若存在变量 export=ExportBuilder,将自动 finalize)
  81. #
  82. # 示例:按 UI 所见优先使用“编辑源码”与 NOHF,导出每个结果的 md/json,同时附带原始与编辑文件
  83. def main(api):
  84. ids = api.get_ids()
  85. eb = api.build_export('custom_export')
  86. for i, rid in enumerate(ids, start=1):
  87. st = api.get_status(rid)
  88. if st['status'] != 'done':
  89. continue
  90. choice = api.choose_texts(rid, prefer_ui=True, prefer_edit=True)
  91. eb.add_text(f'result_{i}_{rid}/content.md', choice['md'] or '')
  92. eb.add_text(f'result_{i}_{rid}/data.json', choice['json'] or '{}')
  93. paths = api.list_paths(rid)
  94. # 附带原始文件
  95. for p in (paths.get('result') or {}).values():
  96. if p and api.path_exists(p):
  97. name = Path(p).name
  98. eb.add_file(p, f'result_{i}_{rid}/raw/{name}')
  99. # 附带编辑文件
  100. for p in (paths.get('edited') or {}).values():
  101. if p and api.path_exists(p):
  102. name = Path(p).name
  103. eb.add_file(p, f'result_{i}_{rid}/edited/{name}')
  104. return eb
  105. """
  106. # ---------------- Helpers ----------------
  107. def read_image_v2(img):
  108. """Read image from URL or local path / PIL.Image. Supports file paths and URLs."""
  109. if isinstance(img, Image.Image):
  110. return img
  111. if isinstance(img, str) and img.startswith(("http://", "https://")):
  112. with requests.get(img, stream=True) as r:
  113. r.raise_for_status()
  114. return Image.open(io.BytesIO(r.content)).convert("RGB")
  115. if isinstance(img, str) and os.path.exists(img):
  116. return Image.open(img).convert("RGB")
  117. try:
  118. img_res = read_image(img, use_native=True)
  119. if isinstance(img_res, tuple) and isinstance(img_res[0], Image.Image):
  120. return img_res[0]
  121. except Exception:
  122. pass
  123. raise ValueError(f"Unsupported image input: {type(img)} / {repr(img)[:200]}")
  124. def create_temp_session_dir():
  125. session_id = uuid.uuid4().hex[:8]
  126. temp_dir = os.path.join(tempfile.gettempdir(), f"dots_ocr_demo_{session_id}")
  127. os.makedirs(temp_dir, exist_ok=True)
  128. return temp_dir, session_id
  129. def classify_parse_failure(exc, min_p, max_p):
  130. """Return a user-friendly error message for known failure causes."""
  131. msg = str(exc)
  132. reasons = []
  133. # Absolute & semantic constraints
  134. if min_p < ABS_MIN_PIXELS:
  135. reasons.append(
  136. f"Min Pixels 过小:{min_p},必须 >= {ABS_MIN_PIXELS}。建议提高 Min Pixels。"
  137. )
  138. if max_p > ABS_MAX_PIXELS:
  139. reasons.append(
  140. f"Max Pixels 过大:{max_p},必须 <= {ABS_MAX_PIXELS}。建议降低 Max Pixels。"
  141. )
  142. if min_p >= max_p:
  143. reasons.append(
  144. f"像素参数不合法:Min Pixels({min_p}) >= Max Pixels({max_p}),必须满足 Min Pixels < Max Pixels。"
  145. )
  146. lower = msg.lower()
  147. if "no results returned from parser" in lower or "no results returned" in lower:
  148. reasons.append(
  149. "解析未返回结果。可能原因:图像过小、Min Pixels 设置过小或过滤过强。"
  150. f"建议:Min Pixels >= {ABS_MIN_PIXELS} 且 Max Pixels <= {ABS_MAX_PIXELS}。"
  151. )
  152. if "failed to read input" in lower or "cannot identify image file" in lower:
  153. reasons.append("无法读取输入文件,请确认文件是否为有效图片或PDF。")
  154. if ("connection" in lower and "refused" in lower) or ("connectionerror" in lower):
  155. reasons.append("无法连接后端推理服务,请检查 Server IP/Port 与服务状态。")
  156. if not reasons:
  157. reasons.append(f"未知错误:{msg}")
  158. detail = "\n".join(f"- {r}" for r in reasons)
  159. cfg = f"(当前参数:min_pixels={min_p}, max_pixels={max_p})"
  160. return f"解析失败:\n{detail}\n{cfg}"
  161. def _is_transient_backend_error(exc: Exception):
  162. lower = str(exc).lower()
  163. # Common signals: connection refused/reset, timeout, gateway, service unavailable
  164. keywords = [
  165. "connection refused",
  166. "connectionerror",
  167. "timeout",
  168. "timed out",
  169. "gateway",
  170. "service unavailable",
  171. "failed to establish a new connection",
  172. "max retries exceeded",
  173. "read timeout",
  174. "connect timeout",
  175. ]
  176. return any(k in lower for k in keywords)
  177. def parse_image_with_high_level_api(parser, image, prompt_mode, fitz_preprocess=False):
  178. """
  179. Calls parser.parse_image with a PIL image (or accepts image path if parser expects path).
  180. Returns dictionary with artifacts. Keeps a temp PNG of the input for traceability.
  181. """
  182. temp_dir, session_id = create_temp_session_dir()
  183. if not isinstance(image, Image.Image):
  184. image = read_image_v2(image)
  185. temp_image_path = os.path.join(temp_dir, f"input_{session_id}.png")
  186. image.save(temp_image_path, "PNG")
  187. filename = f"demo_{session_id}"
  188. results = parser.parse_image(
  189. input_path=image,
  190. filename=filename,
  191. prompt_mode=prompt_mode,
  192. save_dir=temp_dir,
  193. fitz_preprocess=fitz_preprocess,
  194. )
  195. if not results:
  196. raise RuntimeError("No results returned from parser")
  197. result = results[0]
  198. layout_image = None
  199. if result.get("layout_image_path") and os.path.exists(result["layout_image_path"]):
  200. try:
  201. layout_image = Image.open(result["layout_image_path"]).convert("RGB")
  202. except Exception:
  203. layout_image = None
  204. cells_data = None
  205. if result.get("layout_info_path") and os.path.exists(result["layout_info_path"]):
  206. with open(result["layout_info_path"], "r", encoding="utf-8") as f:
  207. cells_data = json.load(f)
  208. md_content = None
  209. if result.get("md_content_path") and os.path.exists(result["md_content_path"]):
  210. with open(result["md_content_path"], "r", encoding="utf-8") as f:
  211. md_content = f.read()
  212. md_content_nohf = None
  213. if result.get("md_content_nohf_path") and os.path.exists(
  214. result["md_content_nohf_path"]
  215. ):
  216. with open(result["md_content_nohf_path"], "r", encoding="utf-8") as f:
  217. md_content_nohf = f.read()
  218. json_code = ""
  219. if cells_data is not None:
  220. try:
  221. json_code = json.dumps(cells_data, ensure_ascii=False, indent=2)
  222. except Exception:
  223. json_code = str(cells_data)
  224. return {
  225. "original_image": image,
  226. "layout_image": layout_image,
  227. "cells_data": cells_data,
  228. "md_content": md_content,
  229. "md_content_nohf": md_content_nohf,
  230. "json_code": json_code,
  231. "filtered": result.get("filtered", False),
  232. "temp_dir": temp_dir,
  233. "session_id": session_id,
  234. "result_paths": result,
  235. "input_width": result.get("input_width", 0),
  236. "input_height": result.get("input_height", 0),
  237. "input_temp_path": temp_image_path,
  238. }
  239. def _validate_pixels(min_p, max_p):
  240. """Coerce pixel parameters. Do NOT auto-swap; semantic errors are handled by pre-validation."""
  241. try:
  242. min_p = int(min_p)
  243. except Exception:
  244. min_p = DEFAULT_CONFIG["min_pixels"]
  245. try:
  246. max_p = int(max_p)
  247. except Exception:
  248. max_p = DEFAULT_CONFIG["max_pixels"]
  249. if min_p <= 0:
  250. min_p = DEFAULT_CONFIG["min_pixels"]
  251. if max_p <= 0:
  252. max_p = DEFAULT_CONFIG["max_pixels"]
  253. return min_p, max_p
  254. def _set_parser_config(server_ip, server_port, min_pixels, max_pixels):
  255. min_pixels, max_pixels = _validate_pixels(min_pixels, max_pixels)
  256. current_config.update(
  257. {
  258. "ip": server_ip,
  259. "port_vllm": int(server_port),
  260. "min_pixels": min_pixels,
  261. "max_pixels": max_pixels,
  262. }
  263. )
  264. dots_parser.ip = server_ip
  265. dots_parser.port = int(server_port)
  266. dots_parser.min_pixels = min_pixels
  267. dots_parser.max_pixels = max_pixels
  268. def purge_queue(rid):
  269. """Best-effort remove tasks matching rid from queue."""
  270. pending = []
  271. try:
  272. while True:
  273. task = TASK_QUEUE.get_nowait()
  274. if task and isinstance(task, tuple):
  275. if task[0] != rid:
  276. pending.append(task)
  277. TASK_QUEUE.task_done()
  278. except queue.Empty:
  279. pass
  280. for t in pending:
  281. TASK_QUEUE.put(t)
  282. # ---------------- Export helpers ----------------
  283. def export_one_rid(rid):
  284. st = RESULTS_CACHE.get(rid)
  285. if not st:
  286. return None
  287. temp_dir = st.get("temp_dir")
  288. if not temp_dir or not os.path.isdir(temp_dir):
  289. return None
  290. out_dir, _sess = create_temp_session_dir()
  291. zip_path = os.path.join(out_dir, f"export_{rid}.zip")
  292. with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
  293. for rt, _, files in os.walk(temp_dir):
  294. for f in files:
  295. src = os.path.join(rt, f)
  296. rel = os.path.relpath(src, temp_dir)
  297. zf.write(src, os.path.join(f"result_{rid}", rel))
  298. return zip_path
  299. def ensure_export_ready(rid):
  300. """Create and cache export zip path if not present."""
  301. st = RESULTS_CACHE.get(rid) or {}
  302. if not st or st.get("status") != "done":
  303. return None
  304. path = st.get("export_path")
  305. if path and os.path.exists(path):
  306. return path
  307. path = export_one_rid(rid)
  308. if path:
  309. st["export_path"] = path
  310. RESULTS_CACHE[rid] = st
  311. return path
  312. # ---------------- Script API & execution ----------------
  313. class ExportBuilder:
  314. def __init__(self, name=None):
  315. root, sid = create_temp_session_dir()
  316. sub = f"script_export_{sid}"
  317. if name:
  318. sub = f"{name}_{sid}"
  319. self.root_dir = os.path.join(root, sub)
  320. os.makedirs(self.root_dir, exist_ok=True)
  321. self._final_zip = None
  322. def _abspath(self, rel_path: str):
  323. rel_path = rel_path.lstrip("/\\")
  324. return os.path.join(self.root_dir, rel_path)
  325. def mkdir(self, rel_dir: str):
  326. p = self._abspath(rel_dir)
  327. os.makedirs(p, exist_ok=True)
  328. return p
  329. def add_text(self, rel_path: str, content: str, encoding: str = "utf-8"):
  330. p = self._abspath(rel_path)
  331. os.makedirs(os.path.dirname(p), exist_ok=True)
  332. with open(p, "w", encoding=encoding) as f:
  333. f.write("" if content is None else str(content))
  334. return p
  335. def add_bytes(self, rel_path: str, data: bytes):
  336. p = self._abspath(rel_path)
  337. os.makedirs(os.path.dirname(p), exist_ok=True)
  338. with open(p, "wb") as f:
  339. f.write(data or b"")
  340. return p
  341. def add_file(self, src_path: str, dest_rel_path: str = None):
  342. if not src_path or not os.path.exists(src_path):
  343. return None
  344. dest_rel_path = dest_rel_path or os.path.basename(src_path)
  345. p = self._abspath(dest_rel_path)
  346. os.makedirs(os.path.dirname(p), exist_ok=True)
  347. shutil.copy2(src_path, p)
  348. return p
  349. def finalize(self, zip_name: str = None):
  350. if self._final_zip and os.path.exists(self._final_zip):
  351. return self._final_zip
  352. out_dir, sid = create_temp_session_dir()
  353. zip_name = zip_name or f"script_export_{sid}.zip"
  354. zip_path = os.path.join(out_dir, zip_name)
  355. with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
  356. for rt, _, files in os.walk(self.root_dir):
  357. for f in files:
  358. src = os.path.join(rt, f)
  359. rel = os.path.relpath(src, self.root_dir)
  360. zf.write(src, rel)
  361. self._final_zip = zip_path
  362. return zip_path
  363. class ScriptAPI:
  364. def __init__(self, ids_snapshot):
  365. self._ids = list(ids_snapshot or [])
  366. def get_ids(self):
  367. return list(self._ids)
  368. def get_status(self, rid: str):
  369. st = dict(RESULTS_CACHE.get(rid) or {})
  370. ui = dict(st.get("ui") or {})
  371. return {
  372. "status": st.get("status", "pending"),
  373. "ui": {
  374. "tab": ui.get("tab", "md"),
  375. "nohf": bool(ui.get("nohf", False)),
  376. "source": ui.get("source", "源码"),
  377. },
  378. "filtered": bool(st.get("filtered", False)),
  379. "input_width": int(st.get("input_width", 0) or 0),
  380. "input_height": int(st.get("input_height", 0) or 0),
  381. }
  382. def get_texts(self, rid: str):
  383. st = dict(RESULTS_CACHE.get(rid) or {})
  384. edits = dict(st.get("edits") or {})
  385. return {
  386. "md": st.get("md_content") or "",
  387. "md_nohf": st.get("md_content_nohf") or "",
  388. "json": st.get("json_code") or "",
  389. "md_edit": edits.get("md"),
  390. "md_nohf_edit": edits.get("nohf"),
  391. "json_edit": edits.get("json"),
  392. }
  393. def choose_texts(
  394. self,
  395. rid: str,
  396. prefer_ui: bool = True,
  397. prefer_edit: bool = True,
  398. prefer_nohf: bool | None = None,
  399. ):
  400. st = dict(RESULTS_CACHE.get(rid) or {})
  401. ui = dict(st.get("ui") or {})
  402. # UI 指示
  403. ui_nohf = bool(ui.get("nohf", False))
  404. ui_source_is_edit = str(ui.get("source", "源码")) == "编辑源码"
  405. # 选择 nohf
  406. use_nohf = ui_nohf if prefer_nohf is None else bool(prefer_nohf)
  407. # 选择是否优先编辑
  408. prefer_edit_final = bool(prefer_edit or (prefer_ui and ui_source_is_edit))
  409. t = self.get_texts(rid)
  410. # Markdown
  411. md_orig = t["md_nohf"] if use_nohf else t["md"]
  412. md_edit = t["md_nohf_edit"] if use_nohf else t["md_edit"]
  413. md = (md_edit if (prefer_edit_final and md_edit is not None) else md_orig) or ""
  414. # JSON
  415. json_text = (
  416. t["json_edit"]
  417. if (prefer_edit_final and t.get("json_edit") is not None)
  418. else t["json"]
  419. ) or ""
  420. return {"md": md, "json": json_text}
  421. def list_paths(self, rid: str):
  422. st = dict(RESULTS_CACHE.get(rid) or {})
  423. rp = dict(st.get("result_paths") or {})
  424. md_p = rp.get("md_content_path")
  425. nohf_p = rp.get("md_content_nohf_path")
  426. json_p = rp.get("layout_info_path") or rp.get("json_path")
  427. image_p = rp.get("layout_image_path") or None
  428. # 编辑路径(若存在)
  429. edited_md = None
  430. edited_nohf = None
  431. edited_json = None
  432. try:
  433. edited_md = _edited_filepath(st, "md")
  434. if not os.path.exists(edited_md):
  435. edited_md = None
  436. except Exception:
  437. edited_md = None
  438. try:
  439. edited_nohf = _edited_filepath(st, "nohf")
  440. if not os.path.exists(edited_nohf):
  441. edited_nohf = None
  442. except Exception:
  443. edited_nohf = None
  444. try:
  445. edited_json = _edited_filepath(st, "json")
  446. if not os.path.exists(edited_json):
  447. edited_json = None
  448. except Exception:
  449. edited_json = None
  450. return {
  451. "temp_dir": st.get("temp_dir"),
  452. "session_id": st.get("session_id"),
  453. "result": {
  454. "md": md_p if (md_p and os.path.exists(md_p)) else None,
  455. "md_nohf": nohf_p if (nohf_p and os.path.exists(nohf_p)) else None,
  456. "json": json_p if (json_p and os.path.exists(json_p)) else None,
  457. "layout": image_p if (image_p and os.path.exists(image_p)) else None,
  458. "input_image": (
  459. st.get("input_temp_path")
  460. if (
  461. st.get("input_temp_path")
  462. and os.path.exists(st.get("input_temp_path"))
  463. )
  464. else None
  465. ),
  466. },
  467. "edited": {
  468. "md": edited_md,
  469. "md_nohf": edited_nohf,
  470. "json": edited_json,
  471. },
  472. }
  473. def path_exists(self, p: str) -> bool:
  474. try:
  475. return bool(p) and os.path.exists(p)
  476. except Exception:
  477. return False
  478. def build_export(self, name: str | None = None):
  479. return ExportBuilder(name=name)
  480. def _safe_builtins():
  481. base = (
  482. __builtins__
  483. if isinstance(__builtins__, dict)
  484. else getattr(__builtins__, "__dict__", {})
  485. )
  486. allow = [
  487. "abs",
  488. "min",
  489. "max",
  490. "sum",
  491. "len",
  492. "range",
  493. "enumerate",
  494. "map",
  495. "filter",
  496. "zip",
  497. "list",
  498. "dict",
  499. "set",
  500. "tuple",
  501. "str",
  502. "int",
  503. "float",
  504. "bool",
  505. "print",
  506. "any",
  507. "all",
  508. "sorted",
  509. ]
  510. return {k: base[k] for k in allow if k in base}
  511. def run_user_script(script_code: str, ids_snapshot):
  512. """
  513. 非流式执行用户脚本,捕获标准输出并返回(zip_path, logs)。
  514. """
  515. api = ScriptAPI(ids_snapshot)
  516. ns = {
  517. "__builtins__": _safe_builtins(),
  518. "api": api,
  519. "json": json,
  520. "re": re,
  521. "math": math,
  522. "datetime": datetime,
  523. "Path": Path,
  524. "io": io,
  525. "ExportBuilder": ExportBuilder,
  526. }
  527. import contextlib
  528. from io import StringIO
  529. buf = StringIO()
  530. zip_path = None
  531. try:
  532. code = script_code or ""
  533. with contextlib.redirect_stdout(buf):
  534. exec(code, ns, ns)
  535. result = None
  536. main_fn = ns.get("main")
  537. if callable(main_fn):
  538. result = main_fn(api)
  539. else:
  540. result = ns.get("RESULT") or ns.get("OUTPUT_PATH")
  541. if isinstance(result, ExportBuilder):
  542. zip_path = result.finalize()
  543. elif isinstance(result, str) and result:
  544. if os.path.isdir(result):
  545. eb = ExportBuilder("script_dir_export")
  546. for rt, _, files in os.walk(result):
  547. for f in files:
  548. src = os.path.join(rt, f)
  549. rel = os.path.relpath(src, result)
  550. eb.add_file(src, rel)
  551. zip_path = eb.finalize()
  552. elif os.path.exists(result):
  553. zip_path = result
  554. if not zip_path:
  555. exp = ns.get("export")
  556. if isinstance(exp, ExportBuilder):
  557. zip_path = exp.finalize()
  558. except Exception as e:
  559. err = f"[Script Error] {type(e).__name__}: {e}"
  560. return None, (buf.getvalue() + "\n" + err)
  561. return (
  562. zip_path if (zip_path and os.path.exists(zip_path)) else None
  563. ), buf.getvalue()
  564. def run_user_script_stream(script_code: str, ids_snapshot):
  565. """生成器:实时输出日志,并在结束时返回下载地址与完成状态。"""
  566. # 日志队列
  567. log_q = queue.Queue()
  568. def _emit(kind, payload=None):
  569. log_q.put((kind, payload))
  570. def debug(*args, **kwargs):
  571. text = " ".join(str(a) for a in args)
  572. if text:
  573. _emit("log", text)
  574. # 准备脚本命名空间(与非流式版本一致,但覆盖 print/debug)
  575. api = ScriptAPI(ids_snapshot)
  576. ns = {
  577. "__builtins__": _safe_builtins(),
  578. "api": api,
  579. "json": json,
  580. "re": re,
  581. "math": math,
  582. "datetime": datetime,
  583. "Path": Path,
  584. "io": io,
  585. "ExportBuilder": ExportBuilder,
  586. # 专用日志函数
  587. "debug": debug,
  588. "print": debug,
  589. }
  590. result_holder = {"zip_path": None, "error": None}
  591. def _worker():
  592. try:
  593. code = script_code or ""
  594. exec(code, ns, ns)
  595. res = None
  596. main_fn = ns.get("main")
  597. if callable(main_fn):
  598. res = main_fn(api)
  599. else:
  600. res = ns.get("RESULT") or ns.get("OUTPUT_PATH")
  601. zip_path = None
  602. if isinstance(res, ExportBuilder):
  603. zip_path = res.finalize()
  604. elif isinstance(res, str) and res:
  605. if os.path.isdir(res):
  606. eb = ExportBuilder("script_dir_export")
  607. for rt, _, files in os.walk(res):
  608. for f in files:
  609. src = os.path.join(rt, f)
  610. rel = os.path.relpath(src, res)
  611. eb.add_file(src, rel)
  612. zip_path = eb.finalize()
  613. elif os.path.exists(res):
  614. zip_path = res
  615. if not zip_path:
  616. exp = ns.get("export")
  617. if isinstance(exp, ExportBuilder):
  618. zip_path = exp.finalize()
  619. result_holder["zip_path"] = (
  620. zip_path if (zip_path and os.path.exists(zip_path)) else None
  621. )
  622. except Exception as e:
  623. result_holder["error"] = f"[Script Error] {type(e).__name__}: {e}"
  624. finally:
  625. _emit("done", None)
  626. # 启动脚本线程
  627. t = threading.Thread(target=_worker, daemon=True)
  628. t.start()
  629. # 初始状态显示
  630. spinner_html = (
  631. "<div style='display:flex;align-items:center;gap:8px;'>"
  632. "<svg width='18' height='18' viewBox='0 0 50 50' style='animation:spin 1s linear infinite'>"
  633. "<circle cx='25' cy='25' r='20' stroke='#FF576D' stroke-width='4' fill='none' stroke-linecap='round' "
  634. "stroke-dasharray='31.4 31.4'>" # dash pattern for arc
  635. "</circle></svg>"
  636. "<style>@keyframes spin{0%{transform:rotate(0deg)}100%{transform:rotate(360deg)}}</style>"
  637. "<span>脚本运行中…</span></div>"
  638. )
  639. log_buf_lines = []
  640. # 初始仅显示运行中动画,日志区域留空
  641. yield None, spinner_html, ""
  642. # 实时拉取日志并渲染
  643. while True:
  644. try:
  645. kind, payload = log_q.get(timeout=0.2)
  646. except queue.Empty:
  647. if not t.is_alive():
  648. # 线程已结束但没有新的事件,跳到收尾
  649. break
  650. else:
  651. continue
  652. if kind == "log":
  653. # 追加日志并推送更新
  654. if isinstance(payload, str):
  655. for line in payload.splitlines() or [payload]:
  656. if line.strip() == "":
  657. continue
  658. log_buf_lines.append(line)
  659. yield None, spinner_html, "```\n" + "\n".join(
  660. log_buf_lines[-200:]
  661. ) + "\n```" # 限制最后200行
  662. elif kind == "done":
  663. break
  664. # 收尾:根据结果/错误输出最终状态
  665. if result_holder.get("error"):
  666. log_buf_lines.append(result_holder["error"])
  667. status_html = (
  668. "<div style='display:flex;align-items:center;gap:8px;color:#fca5a5'>"
  669. "<span>❌ 脚本执行失败</span></div>"
  670. )
  671. yield None, status_html, "```\n" + "\n".join(log_buf_lines[-500:]) + "\n```"
  672. else:
  673. status_html = (
  674. "<div style='display:flex;align-items:center;gap:8px;color:#86efac'>"
  675. "<span>✅ 脚本执行完成</span></div>"
  676. )
  677. if result_holder.get("zip_path"):
  678. yield result_holder["zip_path"], status_html, "```\n" + "\n".join(
  679. log_buf_lines[-500:]
  680. ) + "\n```"
  681. else:
  682. log_buf_lines.append(
  683. "(无可下载文件返回,若需导出请返回 ExportBuilder 或目录/文件路径)"
  684. )
  685. yield None, status_html, "```\n" + "\n".join(log_buf_lines[-500:]) + "\n```"
  686. """
  687. 执行用户脚本,返回 (zip_path or None, log_text)
  688. """
  689. api = ScriptAPI(ids_snapshot)
  690. ns = {
  691. "__builtins__": _safe_builtins(),
  692. "api": api,
  693. # 常用库(只读注入)
  694. "json": json,
  695. "re": re,
  696. "math": math,
  697. "datetime": datetime,
  698. "Path": Path,
  699. "io": io,
  700. # 导出构建器类型(如需构造)
  701. "ExportBuilder": ExportBuilder,
  702. }
  703. import contextlib
  704. from io import StringIO
  705. buf = StringIO()
  706. zip_path = None
  707. try:
  708. code = script_code or ""
  709. with contextlib.redirect_stdout(buf):
  710. exec(code, ns, ns)
  711. result = None
  712. main_fn = ns.get("main")
  713. if callable(main_fn):
  714. result = main_fn(api)
  715. else:
  716. result = ns.get("RESULT") or ns.get("OUTPUT_PATH")
  717. # 结果归档处理
  718. if isinstance(result, ExportBuilder):
  719. zip_path = result.finalize()
  720. elif isinstance(result, str) and result:
  721. if os.path.isdir(result):
  722. eb = ExportBuilder("script_dir_export")
  723. for rt, _, files in os.walk(result):
  724. for f in files:
  725. src = os.path.join(rt, f)
  726. rel = os.path.relpath(src, result)
  727. eb.add_file(src, rel)
  728. zip_path = eb.finalize()
  729. elif os.path.exists(result):
  730. zip_path = result
  731. if not zip_path:
  732. exp = ns.get("export")
  733. if isinstance(exp, ExportBuilder):
  734. zip_path = exp.finalize()
  735. except Exception as e:
  736. err = f"[Script Error] {type(e).__name__}: {e}"
  737. return None, (buf.getvalue() + "\n" + err)
  738. return (
  739. zip_path if (zip_path and os.path.exists(zip_path)) else None
  740. ), buf.getvalue()
  741. def export_selected_rids(ids, selected_labels):
  742. """
  743. Build a combined zip for multiple selected results based on their current images (no reupload).
  744. Only includes items with status == 'done'.
  745. """
  746. if not ids or not selected_labels:
  747. return None
  748. # Map labels "Result N" -> indices
  749. sel_indices = []
  750. for label in selected_labels:
  751. try:
  752. idx = int(str(label).split()[-1]) - 1
  753. if 0 <= idx < len(ids):
  754. sel_indices.append(idx)
  755. except Exception:
  756. continue
  757. if not sel_indices:
  758. return None
  759. out_dir, session_id = create_temp_session_dir()
  760. zip_path = os.path.join(out_dir, f"export_selected_{session_id}.zip")
  761. with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
  762. for i in sel_indices:
  763. rid = ids[i]
  764. st = RESULTS_CACHE.get(rid) or {}
  765. if st.get("status") != "done":
  766. continue
  767. temp_dir = st.get("temp_dir")
  768. if not temp_dir or not os.path.isdir(temp_dir):
  769. # fallback: ensure individual export then include that zip
  770. single_zip = ensure_export_ready(rid)
  771. if single_zip and os.path.exists(single_zip):
  772. zf.write(single_zip, os.path.join(f"result_{i+1}_{rid}.zip"))
  773. continue
  774. base_dir = f"result_{i+1}_{rid}"
  775. for rt, _, files in os.walk(temp_dir):
  776. for f in files:
  777. src = os.path.join(rt, f)
  778. rel = os.path.relpath(src, temp_dir)
  779. zf.write(src, os.path.join(base_dir, rel))
  780. return zip_path if os.path.exists(zip_path) else None
  781. # --------- Edited sources helpers ----------
  782. def _get_base_name_from_result(st: dict):
  783. """Infer base filename like 'demo_xxx' from result paths or session id."""
  784. rp = st.get("result_paths") or {}
  785. for key in ("md_content_path", "md_content_nohf_path", "layout_info_path"):
  786. p = rp.get(key)
  787. if p and isinstance(p, str):
  788. base = os.path.splitext(os.path.basename(p))[0]
  789. if key == "md_content_nohf_path" and base.endswith("_nohf"):
  790. base = base[: -len("_nohf")]
  791. return base
  792. sid = st.get("session_id")
  793. if sid:
  794. return f"demo_{sid}"
  795. return f"demo_{uuid.uuid4().hex[:8]}"
  796. def _edited_dir_for(st: dict):
  797. temp_dir = st.get("temp_dir")
  798. if not temp_dir:
  799. temp_dir, _ = create_temp_session_dir()
  800. st["temp_dir"] = temp_dir
  801. d = os.path.join(temp_dir, "edited")
  802. os.makedirs(d, exist_ok=True)
  803. return d
  804. def _edited_filepath(st: dict, which: str):
  805. """
  806. which in {'md','nohf','json'}
  807. """
  808. base = _get_base_name_from_result(st)
  809. if which == "md":
  810. name = f"{base}.md"
  811. elif which == "nohf":
  812. name = f"{base}_nohf.md"
  813. elif which == "json":
  814. name = f"{base}.json"
  815. else:
  816. raise ValueError(f"unknown edited type: {which}")
  817. return os.path.join(_edited_dir_for(st), name)
  818. def _save_edited_to_disk(st: dict, which: str, content: str):
  819. path = _edited_filepath(st, which)
  820. with open(path, "w", encoding="utf-8") as f:
  821. f.write(content if content is not None else "")
  822. return path
  823. def _delete_edited_from_disk(st: dict, which: str):
  824. try:
  825. path = _edited_filepath(st, which)
  826. if os.path.exists(path):
  827. os.remove(path)
  828. except Exception:
  829. pass
  830. def _invalidate_export_zip(rid: str):
  831. st = RESULTS_CACHE.get(rid) or {}
  832. old = st.get("export_path")
  833. if old and isinstance(old, str) and os.path.exists(old):
  834. try:
  835. os.remove(old)
  836. except Exception:
  837. pass
  838. if "export_path" in st:
  839. st["export_path"] = None
  840. RESULTS_CACHE[rid] = st
  841. # ---------------- UI state helpers (per-card) ----------------
  842. def _default_ui_state():
  843. # 增加 source: '源码' 或 '编辑源码'
  844. return {"preview": True, "nohf": False, "tab": "md", "source": "源码"}
  845. def _ensure_ui_state(rid):
  846. st = RESULTS_CACHE.get(rid) or {}
  847. ui = st.get("ui")
  848. if not isinstance(ui, dict):
  849. ui = _default_ui_state()
  850. st["ui"] = ui
  851. RESULTS_CACHE[rid] = st
  852. else:
  853. # 兼容旧状态缺少新字段
  854. if "source" not in ui:
  855. ui["source"] = "源码"
  856. if "tab" not in ui:
  857. ui["tab"] = "md"
  858. if "preview" not in ui:
  859. ui["preview"] = True
  860. if "nohf" not in ui:
  861. ui["nohf"] = False
  862. RESULTS_CACHE[rid] = st
  863. return ui
  864. # ---------------- Background worker ----------------
  865. def background_processor():
  866. while True:
  867. try:
  868. task = TASK_QUEUE.get(timeout=1)
  869. except queue.Empty:
  870. continue
  871. if task is None:
  872. # Important: mark done for sentinel to keep queue counters balanced
  873. try:
  874. TASK_QUEUE.task_done()
  875. finally:
  876. pass
  877. break
  878. rid, filepath, prompt_mode, server_ip, server_port, min_p, max_p, fitz_flag = (
  879. task
  880. )
  881. image = None
  882. try:
  883. # Build parser instance for this task
  884. local_parser = DotsOCRParser(
  885. ip=server_ip,
  886. port=int(server_port),
  887. dpi=200,
  888. min_pixels=min_p,
  889. max_pixels=max_p,
  890. )
  891. # Read image
  892. try:
  893. fp_lower = str(filepath).lower() if isinstance(filepath, str) else ""
  894. if fitz_flag or fp_lower.endswith(".pdf"):
  895. try:
  896. import fitz as _fitz
  897. doc = _fitz.open(filepath)
  898. page = doc.load_page(0)
  899. pix = page.get_pixmap()
  900. mode = "RGBA" if pix.alpha else "RGB"
  901. image = Image.frombytes(
  902. mode, (pix.width, pix.height), pix.samples
  903. )
  904. doc.close()
  905. except Exception:
  906. image = read_image_v2(filepath)
  907. else:
  908. image = read_image_v2(filepath)
  909. except Exception as e:
  910. raise RuntimeError(f"Failed to read input {filepath}: {e}")
  911. # Parse
  912. result = parse_image_with_high_level_api(
  913. local_parser, image, prompt_mode, fitz_preprocess=fitz_flag
  914. )
  915. result["status"] = "done"
  916. # Preserve source/input path but prefer prev.source_path if available
  917. prev = RESULTS_CACHE.get(rid) or {}
  918. # Preserve UI state across re-parses/results
  919. prev_ui = prev.get("ui") if isinstance(prev, dict) else None
  920. result["ui"] = prev_ui if isinstance(prev_ui, dict) else _default_ui_state()
  921. if isinstance(prev, dict) and isinstance(prev.get("edits"), dict):
  922. result["edits"] = dict(prev.get("edits"))
  923. if isinstance(prev, dict) and prev.get("source_path"):
  924. result["source_path"] = prev.get("source_path")
  925. else:
  926. if isinstance(filepath, str) and os.path.exists(filepath):
  927. result["source_path"] = filepath
  928. else:
  929. result["source_path"] = result.get("input_temp_path")
  930. if isinstance(prev, dict) and prev.get("input_path"):
  931. result["input_path"] = prev.get("input_path")
  932. # Commit result
  933. RESULTS_CACHE[rid] = result
  934. # Pre-build export zip for first-click download
  935. try:
  936. zip_path = ensure_export_ready(rid)
  937. if zip_path:
  938. result = RESULTS_CACHE.get(rid, result)
  939. result["export_path"] = zip_path
  940. RESULTS_CACHE[rid] = result
  941. except Exception:
  942. pass
  943. except Exception as e:
  944. # Auto-retry for transient backend errors (e.g., server down temporarily)
  945. if _is_transient_backend_error(e):
  946. attempts = RETRY_COUNTS.get(rid, 0)
  947. if attempts < MAX_AUTO_RETRIES:
  948. RETRY_COUNTS[rid] = attempts + 1
  949. delay = min(10.0, (RETRY_BACKOFF_BASE**attempts))
  950. # keep state pending, annotate attempts
  951. prev = RESULTS_CACHE.get(rid, {}) or {}
  952. pend_state = dict(prev)
  953. pend_state.update(
  954. {
  955. "status": "pending",
  956. "retry_attempts": attempts + 1,
  957. }
  958. )
  959. RESULTS_CACHE[rid] = pend_state
  960. # Re-enqueue after delay on a timer to avoid blocking worker
  961. def _requeue_later():
  962. TASK_QUEUE.put(
  963. (
  964. rid,
  965. filepath,
  966. prompt_mode,
  967. server_ip,
  968. int(server_port),
  969. min_p,
  970. max_p,
  971. fitz_flag,
  972. )
  973. )
  974. threading.Timer(delay, _requeue_later).start()
  975. # Do not mark error; move on
  976. continue
  977. # Build a rich error state that preserves re-parse materials
  978. prev = RESULTS_CACHE.get(rid, {}) or {}
  979. err_state = dict(prev) # preserve input_path etc.
  980. err_state["status"] = "error"
  981. err_state["md_content"] = classify_parse_failure(e, min_p, max_p)
  982. # Save a temporary PNG for re-parse if we have an image in memory
  983. if isinstance(image, Image.Image):
  984. try:
  985. tmp_dir, _sid = create_temp_session_dir()
  986. tmp_path = os.path.join(tmp_dir, f"error_input_{rid}.png")
  987. image.save(tmp_path, "PNG")
  988. err_state["original_image"] = image
  989. err_state["input_temp_path"] = tmp_path
  990. err_state["temp_dir"] = tmp_dir
  991. except Exception:
  992. err_state["original_image"] = image
  993. if isinstance(filepath, str) and filepath:
  994. err_state.setdefault("source_path", filepath)
  995. # Preserve UI state if missing
  996. if not isinstance(err_state.get("ui"), dict):
  997. err_state["ui"] = _default_ui_state()
  998. RESULTS_CACHE[rid] = err_state
  999. finally:
  1000. # Mark the non-sentinel task as done
  1001. try:
  1002. # If previous branch already marked sentinel done, skip double mark
  1003. if task is not None:
  1004. TASK_QUEUE.task_done()
  1005. except Exception:
  1006. pass
  1007. def _stop_all_workers():
  1008. """Stop all worker threads gracefully by sending sentinels and joining."""
  1009. global WORKER_THREADS
  1010. with THREAD_LOCK:
  1011. n = len(WORKER_THREADS)
  1012. if n == 0:
  1013. return
  1014. # Send one sentinel per worker
  1015. for _ in range(n):
  1016. TASK_QUEUE.put(None)
  1017. # Join all workers
  1018. for t in WORKER_THREADS:
  1019. try:
  1020. t.join(timeout=5.0)
  1021. except Exception:
  1022. pass
  1023. WORKER_THREADS = []
  1024. def _start_workers(count: int):
  1025. """Start exactly `count` worker threads if not already running."""
  1026. global WORKER_THREADS
  1027. with THREAD_LOCK:
  1028. running = len(WORKER_THREADS)
  1029. need = max(0, int(count) - running)
  1030. for _ in range(need):
  1031. t = threading.Thread(target=background_processor, daemon=True)
  1032. t.start()
  1033. WORKER_THREADS.append(t)
  1034. def start_background_processor():
  1035. """Ensure at least one worker is running (used by legacy calls)."""
  1036. _start_workers(max(1, MAX_CONCURRENCY))
  1037. def set_max_concurrency(n: int):
  1038. """Restart worker pool to match desired concurrency."""
  1039. global MAX_CONCURRENCY
  1040. n = int(n) if isinstance(n, (int, float)) else 1
  1041. if n <= 0:
  1042. n = 1
  1043. MAX_CONCURRENCY = n
  1044. # Restart workers to apply new concurrency
  1045. _stop_all_workers()
  1046. _start_workers(MAX_CONCURRENCY)
  1047. # ---------------- Queueing / task helpers ----------------
  1048. def _pixel_reasons(min_p, max_p):
  1049. reasons = []
  1050. if min_p < ABS_MIN_PIXELS:
  1051. reasons.append(f"Min Pixels 过小:{min_p},必须 >= {ABS_MIN_PIXELS}。")
  1052. if max_p > ABS_MAX_PIXELS:
  1053. reasons.append(f"Max Pixels 过大:{max_p},必须 <= {ABS_MAX_PIXELS}。")
  1054. if min_p >= max_p:
  1055. reasons.append(
  1056. f"像素参数不合法:Min Pixels({min_p}) >= Max Pixels({max_p}),必须满足 Min Pixels < Max Pixels。"
  1057. )
  1058. return reasons
  1059. def add_tasks_to_queue(
  1060. file_list, prompt_mode, server_ip, server_port, min_p, max_p, fitz, cur_ids
  1061. ):
  1062. """Queue uploaded file paths (expects file_list of local file paths or tuples (parse_path, source_path))."""
  1063. if not file_list:
  1064. return cur_ids, "No images uploaded."
  1065. min_p, max_p = _validate_pixels(min_p, max_p)
  1066. start_background_processor()
  1067. ids = list(cur_ids or [])
  1068. skipped = 0
  1069. queued = 0
  1070. for fp in file_list:
  1071. # Normalize: support tuple (parse_path, source_path)
  1072. parse_fp = None
  1073. source_fp = None
  1074. if isinstance(fp, (list, tuple)) and len(fp) >= 1:
  1075. parse_fp = fp[0]
  1076. # If tuple contains original source as second element, use it
  1077. source_fp = fp[1] if len(fp) >= 2 else fp[0]
  1078. else:
  1079. parse_fp = fp
  1080. source_fp = fp
  1081. if isinstance(parse_fp, (list, tuple)):
  1082. parse_fp = parse_fp[0] if len(parse_fp) > 0 else None
  1083. rid = uuid.uuid4().hex[:8]
  1084. ids.append(rid)
  1085. # placeholder with input_path so re-parse works even before parse
  1086. RESULTS_CACHE[rid] = {
  1087. "status": "pending",
  1088. "input_path": parse_fp,
  1089. "source_path": source_fp,
  1090. "ui": _default_ui_state(), # 初始化每项的独立 UI 状态
  1091. }
  1092. reason = _pixel_reasons(min_p, max_p)
  1093. if reason:
  1094. RESULTS_CACHE[rid] = {
  1095. "status": "error",
  1096. "md_content": "参数越界,未开始解析:\n"
  1097. + "\n".join(f"- {r}" for r in reason)
  1098. + f"\n(当前参数:min_pixels={min_p}, max_pixels={max_p})",
  1099. "input_path": parse_fp,
  1100. "source_path": source_fp,
  1101. "ui": _default_ui_state(),
  1102. }
  1103. skipped += 1
  1104. continue
  1105. TASK_QUEUE.put(
  1106. (
  1107. rid,
  1108. parse_fp,
  1109. prompt_mode,
  1110. server_ip,
  1111. int(server_port),
  1112. min_p,
  1113. max_p,
  1114. fitz,
  1115. )
  1116. )
  1117. queued += 1
  1118. info = f"Queued {queued} item(s)."
  1119. if skipped:
  1120. info += f" Skipped {skipped} due to invalid pixel limits."
  1121. return ids, info
  1122. def enqueue_single_reparse(
  1123. rid, reupload_path, prompt_mode, server_ip, server_port, min_p, max_p, fitz
  1124. ):
  1125. """
  1126. Enqueue a reparse for single result id.
  1127. Path selection priority:
  1128. reupload_path -> result.source_path -> result.input_temp_path -> result.input_path -> result.original_image (dump to temp PNG)
  1129. """
  1130. min_p, max_p = _validate_pixels(min_p, max_p)
  1131. start_background_processor()
  1132. st = RESULTS_CACHE.get(rid, {}) or {}
  1133. # Pixel constraints: if invalid, set error state and return (do not enqueue)
  1134. reason = _pixel_reasons(min_p, max_p)
  1135. if reason:
  1136. new_state = st.copy()
  1137. new_state.update(
  1138. {
  1139. "status": "error",
  1140. "md_content": "参数越界,未开始解析:\n"
  1141. + "\n".join(f"- {r}" for r in reason)
  1142. + f"\n(当前参数:min_pixels={min_p}, max_pixels={max_p})",
  1143. }
  1144. )
  1145. # 保留 UI 状态
  1146. if "ui" not in new_state:
  1147. new_state["ui"] = _default_ui_state()
  1148. RESULTS_CACHE[rid] = new_state
  1149. return
  1150. if isinstance(reupload_path, (tuple, list)):
  1151. reupload_path = reupload_path[0] if len(reupload_path) > 0 else None
  1152. filepath = None
  1153. if reupload_path:
  1154. filepath = reupload_path
  1155. elif st.get("source_path"):
  1156. filepath = st.get("source_path")
  1157. elif st.get("input_temp_path"):
  1158. filepath = st.get("input_temp_path")
  1159. elif st.get("input_path"):
  1160. filepath = st.get("input_path")
  1161. else:
  1162. img = st.get("original_image")
  1163. if isinstance(img, Image.Image):
  1164. tmp_dir, _ = create_temp_session_dir()
  1165. tmp_path = os.path.join(tmp_dir, f"reparse_{rid}.png")
  1166. try:
  1167. img.save(tmp_path, "PNG")
  1168. filepath = tmp_path
  1169. except Exception:
  1170. filepath = None
  1171. if not filepath:
  1172. new_state = st.copy()
  1173. new_state.update(
  1174. {
  1175. "status": "error",
  1176. "md_content": "重解析失败:未找到可用的图片来源。请重新上传图片或检查缓存目录。",
  1177. }
  1178. )
  1179. if "ui" not in new_state:
  1180. new_state["ui"] = _default_ui_state()
  1181. RESULTS_CACHE[rid] = new_state
  1182. return
  1183. new_state = st.copy()
  1184. new_state.update(
  1185. {
  1186. "status": "pending",
  1187. "input_path": filepath,
  1188. "last_used_config": {
  1189. "ip": server_ip,
  1190. "port": int(server_port),
  1191. "min_pixels": min_p,
  1192. "max_pixels": max_p,
  1193. "prompt_mode": prompt_mode,
  1194. },
  1195. }
  1196. )
  1197. # 保留 UI 状态
  1198. if "ui" not in new_state:
  1199. new_state["ui"] = _default_ui_state()
  1200. RESULTS_CACHE[rid] = new_state
  1201. TASK_QUEUE.put(
  1202. (rid, filepath, prompt_mode, server_ip, int(server_port), min_p, max_p, fitz)
  1203. )
  1204. def delete_one(ids, rid, tick):
  1205. new_ids = [x for x in (ids or []) if x != rid]
  1206. st = RESULTS_CACHE.get(rid)
  1207. temp_dir = st.get("temp_dir") if st else None
  1208. if rid in RESULTS_CACHE:
  1209. del RESULTS_CACHE[rid]
  1210. if rid in RETRY_COUNTS:
  1211. del RETRY_COUNTS[rid]
  1212. purge_queue(rid)
  1213. if temp_dir and os.path.exists(temp_dir):
  1214. threading.Thread(
  1215. target=lambda: shutil.rmtree(temp_dir, ignore_errors=True), daemon=True
  1216. ).start()
  1217. return new_ids, int(tick or 0) + 1
  1218. # ---------------- Gradio UI ----------------
  1219. def create_gradio_interface():
  1220. css = """
  1221. /* basic theme */
  1222. :root { --bg:#0b1220; --card:#111827; --muted:#9ca3af; --accent:#FF576D; --text:#e5e7eb; }
  1223. body, .gradio-container { background: var(--bg) !important; color: var(--text) !important; }
  1224. .result-card { background: var(--card); border:1px solid #1f2937; border-radius:8px; padding:10px; margin-bottom:12px; }
  1225. .muted { color: var(--muted); font-size:0.9em; }
  1226. /* skeleton shimmer */
  1227. .skeleton { position:relative; overflow:hidden; background:#0f172a; border-radius:6px; }
  1228. .skeleton::after {
  1229. content:""; position:absolute; inset:0; transform:translateX(-100%);
  1230. background:linear-gradient(90deg, rgba(255,255,255,0), rgba(255,255,255,0.06), rgba(255,255,255,0));
  1231. animation:shimmer 1.2s infinite;
  1232. }
  1233. @keyframes shimmer { 100% { transform:translateX(100%);} }
  1234. /* Hide unwanted footer/buttons (robust selectors) */
  1235. footer, .footer, #footer, footer[role="contentinfo"] { display:none !important; }
  1236. [aria-label="Use via API"], [aria-label*="API"], [title*="API"], a[href*="/api"], a[href*="api_docs"], a[href*="gradio.app"] { display:none !important; }
  1237. button[aria-label="Settings"], button[aria-label*="设置"], [aria-label="Built with Gradio"] { display:none !important; }
  1238. /* Script log area: single inner scrollbar on <pre>, outer container hidden overflow */
  1239. .script-log { max-height: 260px; overflow: hidden; border:1px solid #1f2937; border-radius:6px; padding:0; }
  1240. .script-log pre {
  1241. max-height: 260px;
  1242. overflow: auto;
  1243. margin: 0;
  1244. padding: 6px;
  1245. background: transparent;
  1246. scrollbar-width: thin; /* Firefox */
  1247. scrollbar-color: rgba(255,255,255,0.2) transparent;
  1248. }
  1249. .script-log pre::-webkit-scrollbar { width: 6px; height: 6px; }
  1250. .script-log pre::-webkit-scrollbar-track { background: transparent; }
  1251. .script-log pre::-webkit-scrollbar-thumb { background: rgba(255,255,255,0.12); border-radius: 4px; }
  1252. .script-log pre:hover::-webkit-scrollbar-thumb { background: rgba(255,255,255,0.25); }
  1253. """
  1254. with gr.Blocks(css=css, title="dots.ocr") as demo:
  1255. # Left column controls
  1256. with gr.Row():
  1257. with gr.Column(scale=1):
  1258. file_input = gr.File(
  1259. label="Upload Multiple Images",
  1260. type="filepath",
  1261. file_count="multiple",
  1262. file_types=[".jpg", ".jpeg", ".png", ".pdf"],
  1263. )
  1264. # Filter out the unwanted 'prompt_grounding_ocr' mode
  1265. allowed_modes = [
  1266. m
  1267. for m in dict_promptmode_to_prompt.keys()
  1268. if m != "prompt_grounding_ocr"
  1269. ]
  1270. if not allowed_modes:
  1271. allowed_modes = list(dict_promptmode_to_prompt.keys())
  1272. prompt_mode = gr.Dropdown(
  1273. label="Prompt Mode",
  1274. choices=allowed_modes,
  1275. value=allowed_modes[0],
  1276. )
  1277. prompt_display = gr.Textbox(
  1278. label="Prompt Preview",
  1279. value=dict_promptmode_to_prompt[allowed_modes[0]],
  1280. interactive=False,
  1281. lines=4,
  1282. )
  1283. with gr.Row():
  1284. parse_btn = gr.Button("🔍 Parse", variant="primary")
  1285. clear_btn = gr.Button("🗑️ Clear")
  1286. with gr.Accordion("Advanced Config", open=False):
  1287. fitz_preprocess = gr.Checkbox(label="fitz_preprocess", value=True)
  1288. server_ip = gr.Textbox(
  1289. label="Server IP", value=DEFAULT_CONFIG["ip"]
  1290. )
  1291. server_port = gr.Number(
  1292. label="Port", value=DEFAULT_CONFIG["port_vllm"], precision=0
  1293. )
  1294. min_pixels = gr.Number(
  1295. label="Min Pixels", value=DEFAULT_CONFIG["min_pixels"]
  1296. )
  1297. max_pixels = gr.Number(
  1298. label="Max Pixels", value=DEFAULT_CONFIG["max_pixels"]
  1299. )
  1300. concurrency = gr.Number(
  1301. label="Max Concurrency",
  1302. value=MAX_CONCURRENCY, # 与实际生效的后台并发保持一致(支持刷新后保持)
  1303. precision=0,
  1304. interactive=True,
  1305. )
  1306. confirm_delete = gr.Checkbox(
  1307. label="删除前确认(推荐)", value=True, interactive=True
  1308. )
  1309. # Right column: results & actions
  1310. with gr.Column(scale=5):
  1311. info_display = gr.Markdown("Waiting...", elem_id="info_box")
  1312. ids_state = gr.State(value=[])
  1313. store_tick = gr.State(value=0)
  1314. render_bump = gr.State(value=0) # 仅用于在状态变化时触发结果重渲染
  1315. confirm_delete_state = gr.State(value=True)
  1316. confirm_delete.change(
  1317. lambda v: v, inputs=[confirm_delete], outputs=[confirm_delete_state]
  1318. )
  1319. progress_timer = gr.Timer(1.0)
  1320. # Actions 面板(多选)
  1321. with gr.Accordion("Actions", open=False):
  1322. selected_group = gr.CheckboxGroup(
  1323. label="Select Items", choices=[], value=[], interactive=True
  1324. )
  1325. with gr.Row():
  1326. select_all_btn = gr.Button("全选")
  1327. clear_sel_btn = gr.Button("清空选择")
  1328. with gr.Row():
  1329. bulk_reparse_btn = gr.Button("🔁 重解析所选")
  1330. delete_selected_btn = gr.Button("🗑️ 删除所选", variant="stop")
  1331. export_selected_btn = gr.DownloadButton("📦 导出所选")
  1332. # 高级脚本导出
  1333. with gr.Accordion("高级脚本", open=False):
  1334. gr.Markdown(
  1335. "在下方编辑并运行自定义 Python 脚本以自由处理当前解析结果并导出为任意目录/文件结构。"
  1336. "<br/>脚本将在受限环境中执行,可通过 api 对象访问只读数据与构建导出压缩包。",
  1337. elem_classes=["muted"],
  1338. )
  1339. script_code = gr.Code(
  1340. label="Python 脚本",
  1341. language="python",
  1342. value=DEFAULT_SCRIPT_TEMPLATE,
  1343. lines=24,
  1344. interactive=True,
  1345. )
  1346. with gr.Row():
  1347. run_script_btn = gr.Button("▶ 运行脚本", variant="primary")
  1348. script_download_btn = gr.DownloadButton("📦 下载脚本输出")
  1349. script_status = gr.HTML("")
  1350. script_log = gr.Markdown(
  1351. "", elem_id="script_log", elem_classes=["script-log"]
  1352. )
  1353. # 流式执行脚本:实时打印日志与运行状态,并在完成后绑定下载按钮
  1354. run_script_btn.click(
  1355. run_user_script_stream,
  1356. inputs=[script_code, ids_state],
  1357. outputs=[script_download_btn, script_status, script_log],
  1358. show_progress="hidden",
  1359. )
  1360. # 批量删除确认面板
  1361. with gr.Row(visible=False) as bulk_delete_confirm_panel:
  1362. gr.Markdown(
  1363. "确认删除所选结果?该操作不可恢复。",
  1364. elem_classes=["muted"],
  1365. )
  1366. bulk_confirm_delete_btn = gr.Button("确认删除", variant="stop")
  1367. bulk_cancel_delete_btn = gr.Button("取消")
  1368. # Render results dynamically
  1369. @gr.render(inputs=[ids_state, render_bump])
  1370. def render_results(ids, _bump):
  1371. if not ids:
  1372. return gr.Markdown("No results yet.")
  1373. with gr.Column():
  1374. for idx, rid in enumerate(ids):
  1375. data = RESULTS_CACHE.get(rid, {}) or {}
  1376. status = data.get("status", "pending")
  1377. # 确保每张卡都有独立 UI 状态(并写回缓存,保证后续使用)
  1378. ui = _ensure_ui_state(rid)
  1379. preview_on = bool(ui.get("preview", True))
  1380. nohf_on = bool(ui.get("nohf", False))
  1381. active_tab = ui.get("tab", "md")
  1382. if active_tab not in ("md", "json"):
  1383. active_tab = "md"
  1384. source_sel = ui.get("source", "源码")
  1385. if source_sel not in ("源码", "编辑源码"):
  1386. source_sel = "源码"
  1387. with gr.Column(
  1388. elem_classes=["result-card"], elem_id=f"card-{rid}"
  1389. ):
  1390. with gr.Row():
  1391. gr.Markdown(
  1392. f"### Result {idx+1} <span class='muted'>RID: {rid}</span>"
  1393. )
  1394. if status == "error":
  1395. gr.Markdown(
  1396. f"⚠️ 解析失败:\n\n{data.get('md_content','Unknown error')}",
  1397. elem_classes=["muted"],
  1398. )
  1399. if status == "done":
  1400. orig_img = data.get("original_image")
  1401. layout_img = data.get("layout_image")
  1402. with gr.Row():
  1403. gr.Image(
  1404. value=orig_img, label="Original", height=300
  1405. )
  1406. gr.Image(
  1407. value=layout_img, label="Layout", height=300
  1408. )
  1409. elif status == "pending":
  1410. with gr.Row():
  1411. gr.HTML(
  1412. "<div class='skeleton' style='width:100%;height:300px;'></div>"
  1413. )
  1414. gr.HTML(
  1415. "<div class='skeleton' style='width:100%;height:300px;'></div>"
  1416. )
  1417. # badges
  1418. with gr.Row():
  1419. badge_md = gr.HTML(
  1420. f"<span class='muted'>MD: {'Preview' if preview_on else 'Source'}</span>"
  1421. )
  1422. badge_nohf = gr.HTML(
  1423. f"<span class='muted'>NOHF: {'On' if nohf_on else 'Off'}</span>"
  1424. )
  1425. # controls
  1426. with gr.Row():
  1427. rid_box = gr.Textbox(value=rid, visible=False)
  1428. preview_cb = gr.Checkbox(
  1429. label="Preview Markdown",
  1430. value=preview_on,
  1431. )
  1432. nohf_cb = gr.Checkbox(label="NOHF", value=nohf_on)
  1433. # 视图切换
  1434. selected_label = (
  1435. "Markdown" if active_tab == "md" else "JSON"
  1436. )
  1437. with gr.Row():
  1438. view_radio = gr.Radio(
  1439. label="视图",
  1440. choices=["Markdown", "JSON"],
  1441. value=selected_label,
  1442. )
  1443. # 内容来源(仅完成状态可用)
  1444. with gr.Row():
  1445. source_radio = gr.Radio(
  1446. label="内容来源",
  1447. choices=["源码", "编辑源码"],
  1448. value=source_sel,
  1449. interactive=True,
  1450. visible=(status == "done"),
  1451. )
  1452. # 内容获取助手
  1453. def _get_texts(rid_value, nohf_flag):
  1454. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1455. md_orig = st.get("md_content") or ""
  1456. md_nohf_orig = st.get("md_content_nohf") or ""
  1457. md_current_orig = (
  1458. md_nohf_orig if nohf_flag else md_orig
  1459. )
  1460. edits = st.get("edits") or {}
  1461. md_edit = (
  1462. edits.get("nohf")
  1463. if nohf_flag
  1464. else edits.get("md")
  1465. )
  1466. if md_edit is None:
  1467. md_edit = md_current_orig
  1468. json_orig = st.get("json_code") or ""
  1469. json_edit = edits.get("json")
  1470. if json_edit is None:
  1471. json_edit = json_orig
  1472. return (
  1473. md_current_orig,
  1474. md_edit,
  1475. json_orig,
  1476. json_edit,
  1477. )
  1478. (
  1479. md_orig_val,
  1480. md_edit_val,
  1481. json_orig_val,
  1482. json_edit_val,
  1483. ) = _get_texts(rid, nohf_on)
  1484. is_md_init = selected_label == "Markdown"
  1485. use_edit_init = source_sel == "编辑源码"
  1486. # 单一预览组件(Markdown 用)
  1487. md_preview = gr.Markdown(
  1488. value=(
  1489. md_edit_val if use_edit_init else md_orig_val
  1490. ),
  1491. visible=(
  1492. status == "done" and is_md_init and preview_on
  1493. ),
  1494. )
  1495. # 原始源码(只读)
  1496. md_code_orig = gr.Code(
  1497. language="markdown",
  1498. value=md_orig_val,
  1499. interactive=False,
  1500. visible=(
  1501. status == "done"
  1502. and is_md_init
  1503. and (not preview_on)
  1504. and (not use_edit_init)
  1505. ),
  1506. )
  1507. # 编辑源码(可编辑、自动保存)
  1508. md_code_edit = gr.Code(
  1509. language="markdown",
  1510. value=md_edit_val,
  1511. interactive=True,
  1512. visible=(
  1513. status == "done"
  1514. and is_md_init
  1515. and (not preview_on)
  1516. and use_edit_init
  1517. ),
  1518. )
  1519. # JSON(原始与编辑)
  1520. json_code_orig = gr.Code(
  1521. language="json",
  1522. value=json_orig_val,
  1523. interactive=False,
  1524. visible=(
  1525. status == "done"
  1526. and (not is_md_init)
  1527. and (not use_edit_init)
  1528. ),
  1529. )
  1530. json_code_edit = gr.Code(
  1531. language="json",
  1532. value=json_edit_val,
  1533. interactive=True,
  1534. visible=(
  1535. status == "done"
  1536. and (not is_md_init)
  1537. and use_edit_init
  1538. ),
  1539. )
  1540. # 仅编辑模式显示
  1541. restore_btn = gr.Button(
  1542. "还原当前内容",
  1543. visible=(status == "done" and use_edit_init),
  1544. )
  1545. # 统一可见性/内容更新
  1546. def _apply_all(
  1547. preview, use_nohf, view_label, src_label, rid_value
  1548. ):
  1549. preview = bool(preview)
  1550. use_nohf = bool(use_nohf)
  1551. is_md = str(view_label) == "Markdown"
  1552. use_edit = str(src_label) == "编辑源码"
  1553. # 写回 UI 状态
  1554. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1555. ui0 = dict(st.get("ui") or _default_ui_state())
  1556. ui0["preview"] = preview
  1557. ui0["nohf"] = use_nohf
  1558. ui0["tab"] = "md" if is_md else "json"
  1559. ui0["source"] = "编辑源码" if use_edit else "源码"
  1560. st["ui"] = ui0
  1561. RESULTS_CACHE[rid_value] = st
  1562. md_o, md_e, j_o, j_e = _get_texts(
  1563. rid_value, use_nohf
  1564. )
  1565. return (
  1566. gr.update(
  1567. value=f"<span class='muted'>MD: {'Preview' if preview else 'Source'}</span>"
  1568. ),
  1569. gr.update(
  1570. value=f"<span class='muted'>NOHF: {'On' if use_nohf else 'Off'}</span>"
  1571. ),
  1572. gr.update(
  1573. value=(md_e if use_edit else md_o),
  1574. visible=(is_md and preview),
  1575. ),
  1576. gr.update(
  1577. value=md_o,
  1578. visible=(
  1579. is_md
  1580. and (not preview)
  1581. and (not use_edit)
  1582. ),
  1583. ),
  1584. gr.update(
  1585. value=md_e,
  1586. visible=(
  1587. is_md and (not preview) and use_edit
  1588. ),
  1589. ),
  1590. gr.update(
  1591. value=j_o,
  1592. visible=(not is_md and (not use_edit)),
  1593. ),
  1594. gr.update(
  1595. value=j_e, visible=(not is_md and use_edit)
  1596. ),
  1597. gr.update(visible=use_edit),
  1598. )
  1599. # 绑定控制项变化:预览、NOHF、视图、来源
  1600. preview_cb.change(
  1601. _apply_all,
  1602. inputs=[
  1603. preview_cb,
  1604. nohf_cb,
  1605. view_radio,
  1606. source_radio,
  1607. rid_box,
  1608. ],
  1609. outputs=[
  1610. badge_md,
  1611. badge_nohf,
  1612. md_preview,
  1613. md_code_orig,
  1614. md_code_edit,
  1615. json_code_orig,
  1616. json_code_edit,
  1617. restore_btn,
  1618. ],
  1619. show_progress="hidden",
  1620. )
  1621. nohf_cb.change(
  1622. _apply_all,
  1623. inputs=[
  1624. preview_cb,
  1625. nohf_cb,
  1626. view_radio,
  1627. source_radio,
  1628. rid_box,
  1629. ],
  1630. outputs=[
  1631. badge_md,
  1632. badge_nohf,
  1633. md_preview,
  1634. md_code_orig,
  1635. md_code_edit,
  1636. json_code_orig,
  1637. json_code_edit,
  1638. restore_btn,
  1639. ],
  1640. show_progress="hidden",
  1641. )
  1642. def _on_view_change(
  1643. view_label,
  1644. rid_value,
  1645. preview_flag,
  1646. nohf_flag,
  1647. src_label,
  1648. ):
  1649. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1650. ui0 = dict(st.get("ui") or _default_ui_state())
  1651. ui0["tab"] = (
  1652. "md"
  1653. if str(view_label) == "Markdown"
  1654. else "json"
  1655. )
  1656. st["ui"] = ui0
  1657. RESULTS_CACHE[rid_value] = st
  1658. return _apply_all(
  1659. preview_flag,
  1660. nohf_flag,
  1661. view_label,
  1662. src_label,
  1663. rid_value,
  1664. )
  1665. view_radio.change(
  1666. _on_view_change,
  1667. inputs=[
  1668. view_radio,
  1669. rid_box,
  1670. preview_cb,
  1671. nohf_cb,
  1672. source_radio,
  1673. ],
  1674. outputs=[
  1675. badge_md,
  1676. badge_nohf,
  1677. md_preview,
  1678. md_code_orig,
  1679. md_code_edit,
  1680. json_code_orig,
  1681. json_code_edit,
  1682. restore_btn,
  1683. ],
  1684. show_progress="hidden",
  1685. )
  1686. def _on_source_change(
  1687. src_label,
  1688. rid_value,
  1689. preview_flag,
  1690. nohf_flag,
  1691. view_label,
  1692. ):
  1693. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1694. ui0 = dict(st.get("ui") or _default_ui_state())
  1695. ui0["source"] = (
  1696. "编辑源码"
  1697. if str(src_label) == "编辑源码"
  1698. else "源码"
  1699. )
  1700. st["ui"] = ui0
  1701. RESULTS_CACHE[rid_value] = st
  1702. return _apply_all(
  1703. preview_flag,
  1704. nohf_flag,
  1705. view_label,
  1706. src_label,
  1707. rid_value,
  1708. )
  1709. source_radio.change(
  1710. _on_source_change,
  1711. inputs=[
  1712. source_radio,
  1713. rid_box,
  1714. preview_cb,
  1715. nohf_cb,
  1716. view_radio,
  1717. ],
  1718. outputs=[
  1719. badge_md,
  1720. badge_nohf,
  1721. md_preview,
  1722. md_code_orig,
  1723. md_code_edit,
  1724. json_code_orig,
  1725. json_code_edit,
  1726. restore_btn,
  1727. ],
  1728. show_progress="hidden",
  1729. )
  1730. # Action buttons per-card
  1731. with gr.Row():
  1732. reparse_btn = gr.Button(
  1733. "🔁 重新解析",
  1734. interactive=(status in ("done", "error")),
  1735. )
  1736. export_btn = gr.DownloadButton(
  1737. "📦 导出",
  1738. interactive=(status == "done"),
  1739. value=(
  1740. data.get("export_path")
  1741. if status == "done"
  1742. else None
  1743. ),
  1744. )
  1745. delete_btn = gr.Button("🗑️ 删除", variant="stop")
  1746. # 自动保存(编辑器变更即写盘 + 刷新导出 + 可能的 Markdown 预览)
  1747. def _save_md_edit(
  1748. val,
  1749. rid_value,
  1750. nohf_flag,
  1751. preview_flag,
  1752. view_label,
  1753. src_label,
  1754. ids,
  1755. selected_labels,
  1756. ):
  1757. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1758. if st.get("status") != "done":
  1759. # 同步“导出所选”以防其它项在编辑(极少见)
  1760. path_sel = export_selected_rids(
  1761. ids, selected_labels
  1762. )
  1763. return (
  1764. gr.update(),
  1765. gr.update(),
  1766. gr.update(value=path_sel),
  1767. )
  1768. which = "nohf" if bool(nohf_flag) else "md"
  1769. edits = dict(st.get("edits") or {})
  1770. edits[which] = val or ""
  1771. st["edits"] = edits
  1772. RESULTS_CACHE[rid_value] = st
  1773. try:
  1774. _save_edited_to_disk(st, which, val or "")
  1775. except Exception:
  1776. pass
  1777. _invalidate_export_zip(rid_value)
  1778. new_zip = ensure_export_ready(rid_value)
  1779. # 刷新“导出所选”
  1780. path_sel = export_selected_rids(
  1781. ids, selected_labels
  1782. )
  1783. # 若当前正处于 Markdown/预览/编辑模式,则更新预览内容
  1784. is_md = str(view_label) == "Markdown"
  1785. use_edit = str(src_label) == "编辑源码"
  1786. if is_md and use_edit and bool(preview_flag):
  1787. return (
  1788. gr.update(value=val or ""),
  1789. gr.update(value=new_zip),
  1790. gr.update(value=path_sel),
  1791. )
  1792. return (
  1793. gr.update(),
  1794. gr.update(value=new_zip),
  1795. gr.update(value=path_sel),
  1796. )
  1797. md_code_edit.change(
  1798. _save_md_edit,
  1799. inputs=[
  1800. md_code_edit,
  1801. rid_box,
  1802. nohf_cb,
  1803. preview_cb,
  1804. view_radio,
  1805. source_radio,
  1806. ids_state,
  1807. selected_group,
  1808. ],
  1809. outputs=[
  1810. md_preview,
  1811. export_btn,
  1812. export_selected_btn,
  1813. ],
  1814. show_progress="hidden",
  1815. )
  1816. def _save_json_edit(
  1817. val, rid_value, ids, selected_labels
  1818. ):
  1819. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1820. if st.get("status") != "done":
  1821. path_sel = export_selected_rids(
  1822. ids, selected_labels
  1823. )
  1824. return gr.update(), gr.update(value=path_sel)
  1825. edits = dict(st.get("edits") or {})
  1826. edits["json"] = val or ""
  1827. st["edits"] = edits
  1828. RESULTS_CACHE[rid_value] = st
  1829. try:
  1830. _save_edited_to_disk(st, "json", val or "")
  1831. except Exception:
  1832. pass
  1833. _invalidate_export_zip(rid_value)
  1834. new_zip = ensure_export_ready(rid_value)
  1835. path_sel = export_selected_rids(
  1836. ids, selected_labels
  1837. )
  1838. return gr.update(value=new_zip), gr.update(
  1839. value=path_sel
  1840. )
  1841. json_code_edit.change(
  1842. _save_json_edit,
  1843. inputs=[
  1844. json_code_edit,
  1845. rid_box,
  1846. ids_state,
  1847. selected_group,
  1848. ],
  1849. outputs=[export_btn, export_selected_btn],
  1850. show_progress="hidden",
  1851. )
  1852. # 还原当前内容
  1853. def _restore_current(
  1854. src_label,
  1855. rid_value,
  1856. nohf_flag,
  1857. preview_flag,
  1858. view_label,
  1859. ids,
  1860. selected_labels,
  1861. ):
  1862. st = RESULTS_CACHE.get(rid_value, {}) or {}
  1863. which = (
  1864. "json"
  1865. if str(view_label) == "JSON"
  1866. else ("nohf" if bool(nohf_flag) else "md")
  1867. )
  1868. # 删除编辑版
  1869. edits = dict(st.get("edits") or {})
  1870. if which in edits:
  1871. edits.pop(which, None)
  1872. st["edits"] = edits
  1873. RESULTS_CACHE[rid_value] = st
  1874. try:
  1875. _delete_edited_from_disk(st, which)
  1876. except Exception:
  1877. pass
  1878. # 重新取原始内容
  1879. md_o, md_e, j_o, j_e = _get_texts(
  1880. rid_value, bool(nohf_flag)
  1881. )
  1882. # 刷新导出
  1883. _invalidate_export_zip(rid_value)
  1884. new_zip = ensure_export_ready(rid_value)
  1885. path_sel = export_selected_rids(
  1886. ids, selected_labels
  1887. )
  1888. # 更新编辑器与预览
  1889. up_md_editor = (
  1890. gr.update(value=md_o)
  1891. if which in ("md", "nohf")
  1892. else gr.update()
  1893. )
  1894. up_json_editor = (
  1895. gr.update(value=j_o)
  1896. if which == "json"
  1897. else gr.update()
  1898. )
  1899. is_md = str(view_label) == "Markdown"
  1900. use_edit = str(src_label) == "编辑源码"
  1901. up_preview = (
  1902. gr.update(value=(md_e if use_edit else md_o))
  1903. if is_md and bool(preview_flag)
  1904. else gr.update()
  1905. )
  1906. return (
  1907. up_md_editor,
  1908. up_json_editor,
  1909. up_preview,
  1910. gr.update(value=new_zip),
  1911. gr.update(value=path_sel),
  1912. )
  1913. restore_btn.click(
  1914. _restore_current,
  1915. inputs=[
  1916. source_radio,
  1917. rid_box,
  1918. nohf_cb,
  1919. preview_cb,
  1920. view_radio,
  1921. ids_state,
  1922. selected_group,
  1923. ],
  1924. outputs=[
  1925. md_code_edit,
  1926. json_code_edit,
  1927. md_preview,
  1928. export_btn,
  1929. export_selected_btn,
  1930. ],
  1931. show_progress="hidden",
  1932. )
  1933. # Reparse panel (collapsed)
  1934. with gr.Column(visible=False) as reparse_panel:
  1935. gr.Markdown("**重解析**")
  1936. with gr.Row():
  1937. reparse_current_btn = gr.Button(
  1938. "基于当前图片直接重解析", variant="primary"
  1939. )
  1940. # Delete confirm panel (collapsed)
  1941. with gr.Row(visible=False) as delete_confirm_panel:
  1942. gr.Markdown(
  1943. "确认删除该结果?该操作不可恢复。",
  1944. elem_classes=["muted"],
  1945. )
  1946. confirm_delete_btn = gr.Button(
  1947. "确认删除", variant="stop"
  1948. )
  1949. cancel_delete_btn = gr.Button("取消")
  1950. # 绑定其他交互
  1951. reparse_btn.click(
  1952. lambda: gr.update(visible=True),
  1953. outputs=[reparse_panel],
  1954. show_progress="hidden",
  1955. )
  1956. def _start_reparse_current(
  1957. rid_value,
  1958. p_mode,
  1959. ip_addr,
  1960. port_val,
  1961. minp,
  1962. maxp,
  1963. fitz_flag,
  1964. tick,
  1965. ids,
  1966. selected_labels,
  1967. ):
  1968. try:
  1969. enqueue_single_reparse(
  1970. rid_value,
  1971. None,
  1972. p_mode,
  1973. ip_addr,
  1974. int(port_val),
  1975. int(minp),
  1976. int(maxp),
  1977. fitz_flag,
  1978. )
  1979. # 重建“导出所选”
  1980. path_sel = export_selected_rids(
  1981. ids, selected_labels
  1982. )
  1983. return (
  1984. int(tick or 0) + 1,
  1985. gr.update(visible=False),
  1986. gr.update(value=path_sel),
  1987. )
  1988. except Exception as e:
  1989. RESULTS_CACHE[rid_value] = {
  1990. "status": "error",
  1991. "md_content": f"Reparse error: {e}",
  1992. # 保留 UI 状态
  1993. "ui": _ensure_ui_state(rid_value),
  1994. }
  1995. path_sel = export_selected_rids(
  1996. ids, selected_labels
  1997. )
  1998. return (
  1999. int(tick or 0) + 1,
  2000. gr.update(visible=False),
  2001. gr.update(value=path_sel),
  2002. )
  2003. reparse_current_btn.click(
  2004. _start_reparse_current,
  2005. inputs=[
  2006. rid_box,
  2007. prompt_mode,
  2008. server_ip,
  2009. server_port,
  2010. min_pixels,
  2011. max_pixels,
  2012. fitz_preprocess,
  2013. store_tick,
  2014. ids_state,
  2015. selected_group,
  2016. ],
  2017. outputs=[
  2018. store_tick,
  2019. reparse_panel,
  2020. export_selected_btn,
  2021. ],
  2022. show_progress="hidden",
  2023. )
  2024. def _on_delete_click(
  2025. rid_value, ids, need_confirm, tick
  2026. ):
  2027. # 如果需要确认,仅展开确认面板,不修改选择框/导出按钮
  2028. if need_confirm:
  2029. return (
  2030. gr.update(visible=True),
  2031. ids,
  2032. tick,
  2033. gr.update(), # selected_group 不变
  2034. gr.update(), # export button 不变
  2035. )
  2036. # 直接删除:更新 ids/tick,并同步 Actions 的选择项与导出按钮
  2037. new_ids, new_tick = delete_one(ids, rid_value, tick)
  2038. choices = [
  2039. f"Result {i+1}"
  2040. for i in range(len(new_ids or []))
  2041. ]
  2042. return (
  2043. gr.update(visible=False),
  2044. new_ids,
  2045. new_tick,
  2046. gr.update(choices=choices, value=[]),
  2047. gr.update(value=None), # 清空导出
  2048. )
  2049. # 单卡删除输出同步 selected_group 与 export_selected_btn
  2050. delete_btn.click(
  2051. _on_delete_click,
  2052. inputs=[
  2053. rid_box,
  2054. ids_state,
  2055. confirm_delete_state,
  2056. store_tick,
  2057. ],
  2058. outputs=[
  2059. delete_confirm_panel,
  2060. ids_state,
  2061. store_tick,
  2062. selected_group,
  2063. export_selected_btn,
  2064. ],
  2065. show_progress="hidden",
  2066. )
  2067. def _confirm_delete(rid_value, ids, tick):
  2068. new_ids, new_tick = delete_one(ids, rid_value, tick)
  2069. choices = [
  2070. f"Result {i+1}"
  2071. for i in range(len(new_ids or []))
  2072. ]
  2073. return (
  2074. new_ids,
  2075. new_tick,
  2076. gr.update(visible=False),
  2077. gr.update(choices=choices, value=[]),
  2078. gr.update(value=None),
  2079. )
  2080. # 确认删除后同步 selected_group 与 export_selected_btn
  2081. confirm_delete_btn.click(
  2082. _confirm_delete,
  2083. inputs=[rid_box, ids_state, store_tick],
  2084. outputs=[
  2085. ids_state,
  2086. store_tick,
  2087. delete_confirm_panel,
  2088. selected_group,
  2089. export_selected_btn,
  2090. ],
  2091. show_progress="hidden",
  2092. )
  2093. cancel_delete_btn.click(
  2094. lambda: gr.update(visible=False),
  2095. outputs=[delete_confirm_panel],
  2096. show_progress="hidden",
  2097. )
  2098. # Top-level events
  2099. def _on_prompt_mode_change(m):
  2100. return dict_promptmode_to_prompt.get(m, "")
  2101. prompt_mode.change(
  2102. fn=_on_prompt_mode_change,
  2103. inputs=[prompt_mode],
  2104. outputs=[prompt_display],
  2105. show_progress="hidden",
  2106. )
  2107. def process_images_simple(
  2108. file_list,
  2109. p_mode,
  2110. server_ip_val,
  2111. server_port_val,
  2112. min_p_val,
  2113. max_p_val,
  2114. fitz_val,
  2115. cur_ids,
  2116. tick,
  2117. ):
  2118. """
  2119. Process images with selected prompt mode. Grounding mode is removed; all files go through normal path.
  2120. """
  2121. minp, maxp = _validate_pixels(min_p_val, max_p_val)
  2122. _set_parser_config(server_ip_val, server_port_val, minp, maxp)
  2123. # normalize file_list (gradio file element may pass nested lists)
  2124. files = []
  2125. if not file_list:
  2126. return (
  2127. gr.update(value=None),
  2128. gr.update(value="No files uploaded."),
  2129. cur_ids,
  2130. tick,
  2131. gr.update(choices=[], value=[]),
  2132. gr.update(value=None), # 清空导出
  2133. )
  2134. # build normalized list
  2135. for f in file_list:
  2136. if isinstance(f, (list, tuple)):
  2137. files.append(f[0] if len(f) > 0 else None)
  2138. else:
  2139. files.append(f)
  2140. # Normal path: queue originals
  2141. new_ids, info = add_tasks_to_queue(
  2142. files,
  2143. p_mode,
  2144. server_ip_val,
  2145. server_port_val,
  2146. minp,
  2147. maxp,
  2148. fitz_val,
  2149. cur_ids,
  2150. )
  2151. # Update checkbox group choices
  2152. choices = [f"Result {i+1}" for i in range(len(new_ids or []))]
  2153. return (
  2154. gr.update(value=None),
  2155. gr.update(value=info),
  2156. new_ids,
  2157. int(tick or 0) + 1,
  2158. gr.update(choices=choices, value=[]),
  2159. gr.update(value=None), # 清空导出
  2160. )
  2161. parse_btn.click(
  2162. fn=process_images_simple,
  2163. inputs=[
  2164. file_input,
  2165. prompt_mode,
  2166. server_ip,
  2167. server_port,
  2168. min_pixels,
  2169. max_pixels,
  2170. fitz_preprocess,
  2171. ids_state,
  2172. store_tick,
  2173. ],
  2174. outputs=[
  2175. file_input,
  2176. info_display,
  2177. ids_state,
  2178. store_tick,
  2179. selected_group,
  2180. export_selected_btn,
  2181. ],
  2182. show_progress="hidden",
  2183. )
  2184. # Concurrency change handler: apply immediately
  2185. def _on_concurrency_change(n):
  2186. try:
  2187. set_max_concurrency(int(n))
  2188. return gr.update(value=f"并发已设置为 {int(n)}。")
  2189. except Exception as e:
  2190. return gr.update(value=f"设置并发失败:{e}")
  2191. concurrency.change(
  2192. _on_concurrency_change,
  2193. inputs=[concurrency],
  2194. outputs=[info_display],
  2195. show_progress="hidden",
  2196. )
  2197. # 会话加载时同步 UI 与当前真实并发(解决刷新后 UI 值与实际不一致)
  2198. def _sync_concurrency_on_session_load():
  2199. try:
  2200. # 如有需要,补齐 worker 到目标并发数(不会减少已有线程)
  2201. _start_workers(max(1, MAX_CONCURRENCY))
  2202. return (
  2203. gr.update(value=int(MAX_CONCURRENCY)),
  2204. gr.update(
  2205. value=f"已同步当前并发为 {int(MAX_CONCURRENCY)}。"
  2206. ),
  2207. )
  2208. except Exception as e:
  2209. return (
  2210. gr.update(value=int(MAX_CONCURRENCY)),
  2211. gr.update(value=f"同步并发时发生异常:{e}"),
  2212. )
  2213. demo.load(
  2214. _sync_concurrency_on_session_load,
  2215. inputs=None,
  2216. outputs=[concurrency, info_display],
  2217. )
  2218. # 生成导出 ZIP(基于当前选择),用于首次点击即可下载
  2219. def _update_export_for_selection(ids, selected_labels):
  2220. path = export_selected_rids(ids, selected_labels)
  2221. return gr.update(
  2222. value=path if path and os.path.exists(path) else None
  2223. )
  2224. # Actions: 全选/清空
  2225. def _select_all(ids):
  2226. choices = [f"Result {i+1}" for i in range(len(ids or []))]
  2227. # 预生成 zip
  2228. path = export_selected_rids(ids, choices)
  2229. return (
  2230. gr.update(choices=choices, value=choices),
  2231. gr.update(
  2232. value=path if path and os.path.exists(path) else None
  2233. ),
  2234. )
  2235. def _clear_selection(ids):
  2236. choices = [f"Result {i+1}" for i in range(len(ids or []))]
  2237. return (
  2238. gr.update(choices=choices, value=[]),
  2239. gr.update(value=None),
  2240. )
  2241. select_all_btn.click(
  2242. _select_all,
  2243. inputs=[ids_state],
  2244. outputs=[selected_group, export_selected_btn],
  2245. show_progress="hidden",
  2246. )
  2247. clear_sel_btn.click(
  2248. _clear_selection,
  2249. inputs=[ids_state],
  2250. outputs=[selected_group, export_selected_btn],
  2251. show_progress="hidden",
  2252. )
  2253. # 当用户手动变更选择时,预构建导出 zip 并绑定到按钮
  2254. selected_group.change(
  2255. _update_export_for_selection,
  2256. inputs=[ids_state, selected_group],
  2257. outputs=[export_selected_btn],
  2258. show_progress="hidden",
  2259. )
  2260. # Actions: 批量重解析(基于当前图片)
  2261. def bulk_reparse(
  2262. selected_labels, ids, p_mode, ip, port, minp, maxp, fitz, tick
  2263. ):
  2264. if not ids or not selected_labels:
  2265. path_sel = export_selected_rids(ids, selected_labels)
  2266. return (
  2267. gr.update(value="未选择任何结果。"),
  2268. int(tick or 0),
  2269. gr.update(value=path_sel),
  2270. )
  2271. # Map labels -> rids
  2272. count = 0
  2273. for label in selected_labels:
  2274. try:
  2275. idx = int(str(label).split()[-1]) - 1
  2276. rid = ids[idx]
  2277. enqueue_single_reparse(
  2278. rid,
  2279. None,
  2280. p_mode,
  2281. ip,
  2282. int(port),
  2283. int(minp),
  2284. int(maxp),
  2285. fitz,
  2286. )
  2287. count += 1
  2288. except Exception:
  2289. continue
  2290. path_sel = export_selected_rids(ids, selected_labels)
  2291. return (
  2292. gr.update(value=f"已触发 {count} 个重解析任务。"),
  2293. int(tick or 0) + 1,
  2294. gr.update(value=path_sel),
  2295. )
  2296. bulk_reparse_btn.click(
  2297. bulk_reparse,
  2298. inputs=[
  2299. selected_group,
  2300. ids_state,
  2301. prompt_mode,
  2302. server_ip,
  2303. server_port,
  2304. min_pixels,
  2305. max_pixels,
  2306. fitz_preprocess,
  2307. store_tick,
  2308. ],
  2309. outputs=[info_display, store_tick, export_selected_btn],
  2310. show_progress="hidden",
  2311. )
  2312. # Actions: 删除所选(尊重“删除前确认”)
  2313. def delete_selected_action(ids, selected_labels, tick):
  2314. # 先从“原始 ids 列表”解析出要删除的 rid 列表,避免索引随删除而错位
  2315. if not ids or not selected_labels:
  2316. choices = [f"Result {i+1}" for i in range(len(ids or []))]
  2317. return (
  2318. ids,
  2319. int(tick or 0),
  2320. gr.update(choices=choices, value=[]),
  2321. gr.update(value=None),
  2322. )
  2323. # 解析 label -> index(去重、过滤非法)
  2324. sel_indices = []
  2325. for label in selected_labels:
  2326. try:
  2327. idx = int(str(label).split()[-1]) - 1
  2328. if 0 <= idx < len(ids):
  2329. sel_indices.append(idx)
  2330. except Exception:
  2331. continue
  2332. if not sel_indices:
  2333. choices = [f"Result {i+1}" for i in range(len(ids or []))]
  2334. return (
  2335. ids,
  2336. int(tick or 0),
  2337. gr.update(choices=choices, value=[]),
  2338. gr.update(value=None),
  2339. )
  2340. sel_indices = sorted(set(sel_indices))
  2341. rids_to_delete = [ids[i] for i in sel_indices]
  2342. new_ids = list(ids)
  2343. new_tick = int(tick or 0)
  2344. # 基于 rid 删除,避免受索引变化影响
  2345. for rid in rids_to_delete:
  2346. new_ids, new_tick = delete_one(new_ids, rid, new_tick)
  2347. choices = [f"Result {i+1}" for i in range(len(new_ids or []))]
  2348. return (
  2349. new_ids,
  2350. new_tick,
  2351. gr.update(choices=choices, value=[]),
  2352. gr.update(value=None),
  2353. )
  2354. def _on_bulk_delete_click(ids, selected_labels, need_confirm, tick):
  2355. if need_confirm:
  2356. # 展示确认面板,不改动任何选择与导出
  2357. return (
  2358. gr.update(visible=True),
  2359. ids,
  2360. tick,
  2361. gr.update(),
  2362. gr.update(),
  2363. )
  2364. # 直接删除并隐藏确认面板
  2365. new_ids, new_tick, sel_update, export_update = (
  2366. delete_selected_action(ids, selected_labels, tick)
  2367. )
  2368. return (
  2369. gr.update(visible=False),
  2370. new_ids,
  2371. new_tick,
  2372. sel_update,
  2373. export_update,
  2374. )
  2375. delete_selected_btn.click(
  2376. _on_bulk_delete_click,
  2377. inputs=[
  2378. ids_state,
  2379. selected_group,
  2380. confirm_delete_state,
  2381. store_tick,
  2382. ],
  2383. outputs=[
  2384. bulk_delete_confirm_panel,
  2385. ids_state,
  2386. store_tick,
  2387. selected_group,
  2388. export_selected_btn,
  2389. ],
  2390. show_progress="hidden",
  2391. )
  2392. def _bulk_confirm_delete(ids, selected_labels, tick):
  2393. new_ids, new_tick, sel_update, export_update = (
  2394. delete_selected_action(ids, selected_labels, tick)
  2395. )
  2396. return (
  2397. new_ids,
  2398. new_tick,
  2399. sel_update,
  2400. export_update,
  2401. gr.update(visible=False),
  2402. )
  2403. bulk_confirm_delete_btn.click(
  2404. _bulk_confirm_delete,
  2405. inputs=[ids_state, selected_group, store_tick],
  2406. outputs=[
  2407. ids_state,
  2408. store_tick,
  2409. selected_group,
  2410. export_selected_btn,
  2411. bulk_delete_confirm_panel,
  2412. ],
  2413. show_progress="hidden",
  2414. )
  2415. bulk_cancel_delete_btn.click(
  2416. lambda: gr.update(visible=False),
  2417. outputs=[bulk_delete_confirm_panel],
  2418. show_progress="hidden",
  2419. )
  2420. # 进度信息
  2421. def update_progress_info(ids, tick, bump):
  2422. if not ids:
  2423. return (
  2424. gr.update(value="Waiting..."),
  2425. tick,
  2426. int(bump or 0),
  2427. )
  2428. pending = 0
  2429. done = 0
  2430. errors = 0
  2431. status_signature = []
  2432. for rid in ids:
  2433. st = RESULTS_CACHE.get(rid, {})
  2434. status = st.get("status", "pending")
  2435. status_signature.append((rid, status))
  2436. if status == "done":
  2437. done += 1
  2438. elif status == "error":
  2439. errors += 1
  2440. else:
  2441. pending += 1
  2442. qsize = TASK_QUEUE.qsize()
  2443. running = max(0, pending - qsize)
  2444. # Info text
  2445. if pending == 0:
  2446. info = (
  2447. f"进度:完成 {done}"
  2448. + ("" if errors == 0 else f",错误 {errors}")
  2449. + "。"
  2450. )
  2451. else:
  2452. info = f"进度:完成 {done},错误 {errors},正在解析 {running},排队 {qsize},待处理合计 {pending}。"
  2453. # Only bump render when any item's status changed
  2454. sig_tuple = tuple(status_signature)
  2455. last_sig = getattr(update_progress_info, "_last_status_sig", None)
  2456. bump_out = int(bump or 0)
  2457. if last_sig != sig_tuple:
  2458. setattr(update_progress_info, "_last_status_sig", sig_tuple)
  2459. bump_out = bump_out + 1
  2460. # Only tick when coarse counts change (avoid unnecessary churn)
  2461. key = f"{done}_{errors}_{pending}"
  2462. last_key = getattr(update_progress_info, "_last_counts_key", None)
  2463. new_tick = int(tick or 0)
  2464. if last_key != key:
  2465. setattr(update_progress_info, "_last_counts_key", key)
  2466. new_tick = new_tick + 1
  2467. return (
  2468. gr.update(value=info),
  2469. new_tick,
  2470. bump_out,
  2471. )
  2472. # 计时器不再触达 selected_group,杜绝与用户交互竞争导致选择重置/计时停止
  2473. progress_timer.tick(
  2474. fn=update_progress_info,
  2475. inputs=[ids_state, store_tick, render_bump],
  2476. outputs=[info_display, store_tick, render_bump],
  2477. show_progress="hidden",
  2478. )
  2479. # Clear all
  2480. def clear_all():
  2481. global RESULTS_CACHE
  2482. while not TASK_QUEUE.empty():
  2483. try:
  2484. TASK_QUEUE.get_nowait()
  2485. TASK_QUEUE.task_done()
  2486. except queue.Empty:
  2487. break
  2488. RESULTS_CACHE = {}
  2489. RETRY_COUNTS.clear()
  2490. # Do not stop workers; keep them alive
  2491. return (
  2492. [],
  2493. 0,
  2494. gr.update(value="Waiting..."),
  2495. 0,
  2496. gr.update(choices=[], value=[]),
  2497. gr.update(value=None),
  2498. )
  2499. clear_btn.click(
  2500. clear_all,
  2501. inputs=None,
  2502. outputs=[
  2503. ids_state,
  2504. store_tick,
  2505. info_display,
  2506. render_bump,
  2507. selected_group,
  2508. export_selected_btn,
  2509. ],
  2510. show_progress="hidden",
  2511. )
  2512. return demo
  2513. # ---------------- main ----------------
  2514. def _queue_compat(blocks: gr.Blocks):
  2515. """
  2516. Gradio version compatibility layer for Blocks.queue:
  2517. - Try Gradio 4.x: default_concurrency_limit + status_update_rate
  2518. - Fallback to Gradio 3.x: concurrency_count + status_update_rate
  2519. - Final fallback: no-arg queue()
  2520. """
  2521. try:
  2522. # Gradio 4.x path
  2523. return blocks.queue(default_concurrency_limit=20, status_update_rate=0.2)
  2524. except TypeError:
  2525. try:
  2526. # Gradio 3.x path
  2527. return blocks.queue(concurrency_count=16, status_update_rate=0.2)
  2528. except TypeError:
  2529. # Minimal fallback
  2530. return blocks.queue()
  2531. def _launch_compat(app: gr.Blocks, port: int):
  2532. """
  2533. Gradio version compatibility for launch parameters.
  2534. """
  2535. try:
  2536. app.launch(
  2537. server_name="0.0.0.0",
  2538. server_port=port,
  2539. debug=True,
  2540. show_api=False, # 3.x/部分4.x可用
  2541. )
  2542. except TypeError:
  2543. # Fallback without show_api
  2544. app.launch(
  2545. server_name="0.0.0.0",
  2546. server_port=port,
  2547. debug=True,
  2548. )
  2549. if __name__ == "__main__":
  2550. import sys
  2551. port = int(sys.argv[1]) if len(sys.argv) > 1 else 7860
  2552. demo = create_gradio_interface()
  2553. app = _queue_compat(demo)
  2554. _launch_compat(app, port)