343 lines
13 KiB
Python
343 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import io
|
|
import json
|
|
import mimetypes
|
|
import threading
|
|
import traceback
|
|
import zipfile
|
|
from http import HTTPStatus
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from itertools import count
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import parse_qs, unquote, urlparse
|
|
|
|
import pipeline
|
|
|
|
|
|
STATIC_DIR = Path(__file__).resolve().parent / "static"
|
|
JOBS: dict[str, dict[str, Any]] = {}
|
|
JOB_IDS = count(1)
|
|
JOB_LOCK = threading.Lock()
|
|
|
|
|
|
def _json_bytes(payload: Any) -> bytes:
|
|
return json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
|
|
|
|
|
|
def _new_job(kind: str, target, *args, **kwargs) -> dict[str, Any]:
|
|
job_id = str(next(JOB_IDS))
|
|
job = {
|
|
"id": job_id,
|
|
"kind": kind,
|
|
"status": "running",
|
|
"logs": [],
|
|
"result": None,
|
|
"error": None,
|
|
}
|
|
with JOB_LOCK:
|
|
JOBS[job_id] = job
|
|
|
|
def log(message: str) -> None:
|
|
with JOB_LOCK:
|
|
job["logs"].append(message)
|
|
job["logs"] = job["logs"][-400:]
|
|
|
|
def wrapped() -> None:
|
|
try:
|
|
job["result"] = target(log, *args, **kwargs)
|
|
job["status"] = "done"
|
|
except Exception as exc:
|
|
job["status"] = "error"
|
|
job["error"] = str(exc)
|
|
log(traceback.format_exc())
|
|
|
|
threading.Thread(target=wrapped, daemon=True).start()
|
|
return job
|
|
|
|
|
|
def _run_dehaze_job(log, payload: dict[str, Any]) -> dict[str, Any]:
|
|
filename = payload["image"]
|
|
methods = payload.get("methods") or []
|
|
options = payload.get("options") or {}
|
|
postprocessors = payload.get("postprocessors") or []
|
|
post_sources = payload.get("post_sources") or []
|
|
reference = payload.get("reference") or filename
|
|
|
|
completed = []
|
|
failed = []
|
|
for method in methods:
|
|
try:
|
|
result = pipeline.run_dehaze_method(filename, method, options.get(method, options), log)
|
|
completed.append({"method": method, "path": pipeline.relpath(result)})
|
|
except Exception as exc:
|
|
failed.append({"method": method, "error": str(exc)})
|
|
log(f"[{method}] 失败: {exc}")
|
|
|
|
if postprocessors:
|
|
sources = post_sources or [item["method"] for item in completed]
|
|
for source in sources:
|
|
try:
|
|
pipeline.run_postprocessors_for_source(
|
|
filename,
|
|
source,
|
|
postprocessors,
|
|
params=payload.get("post_params") or {},
|
|
reference_filename=reference,
|
|
log=log,
|
|
)
|
|
except Exception as exc:
|
|
failed.append({"postprocess_source": source, "error": str(exc)})
|
|
log(f"[后处理/{source}] 失败: {exc}")
|
|
|
|
return {"completed": completed, "failed": failed, "results": pipeline.get_results(filename)}
|
|
|
|
|
|
def _run_post_job(log, payload: dict[str, Any]) -> dict[str, Any]:
|
|
filename = payload["image"]
|
|
source = payload.get("source") or "DCP"
|
|
processors = payload.get("processors") or []
|
|
reference = payload.get("reference") or filename
|
|
outputs = pipeline.run_postprocessors_for_source(
|
|
filename,
|
|
source,
|
|
processors,
|
|
params=payload.get("params") or {},
|
|
reference_filename=reference,
|
|
log=log,
|
|
)
|
|
return {
|
|
"outputs": [pipeline.relpath(path) for path in outputs],
|
|
"results": pipeline.get_results(filename),
|
|
}
|
|
|
|
|
|
def _run_default_batch_job(log, payload: dict[str, Any]) -> dict[str, Any]:
|
|
requested = payload.get("images") or [item["name"] for item in pipeline.list_images()]
|
|
skip_existing = bool(payload.get("skip_existing", True))
|
|
completed = []
|
|
failed = []
|
|
|
|
for filename in requested:
|
|
try:
|
|
log(f"[默认流程] {filename}: Baidu_API -> 自动 S/V")
|
|
baidu_output = pipeline.dehaze_result_path(filename, "Baidu_API")
|
|
if skip_existing and baidu_output.exists():
|
|
log(f"[Baidu_API] 已存在，跳过: {pipeline.relpath(baidu_output)}")
|
|
else:
|
|
baidu_output = pipeline.run_dehaze_method(filename, "Baidu_API", {}, log)
|
|
|
|
auto_output = pipeline.post_result_path(filename, "Baidu_API", "auto_sv", {})
|
|
auto_meta = pipeline.post_metadata_path(auto_output)
|
|
if skip_existing and auto_output.exists() and auto_meta.exists():
|
|
log(f"[自动 S/V] 已存在，跳过: {pipeline.relpath(auto_output)}")
|
|
else:
|
|
if auto_output.exists() and not auto_meta.exists():
|
|
log("[自动 S/V] 已有图片但缺少 S/V 元数据，重新计算")
|
|
outputs = pipeline.run_postprocessors_for_source(
|
|
filename,
|
|
"Baidu_API",
|
|
["auto_sv"],
|
|
params={"auto_sv": {}},
|
|
reference_filename=filename,
|
|
log=log,
|
|
)
|
|
auto_output = outputs[-1]
|
|
|
|
completed.append(
|
|
{
|
|
"image": filename,
|
|
"baidu": pipeline.relpath(baidu_output),
|
|
"auto_sv": pipeline.relpath(auto_output),
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
failed.append({"image": filename, "error": str(exc)})
|
|
log(f"[默认流程] {filename} 失败: {exc}")
|
|
|
|
return {
|
|
"completed": completed,
|
|
"failed": failed,
|
|
"results": {item["image"]: pipeline.get_results(item["image"]) for item in completed},
|
|
}
|
|
|
|
|
|
class DehazeRequestHandler(BaseHTTPRequestHandler):
|
|
server_version = "DehazeWeb/1.0"
|
|
|
|
def log_message(self, fmt: str, *args) -> None:
|
|
print("[%s] %s" % (self.log_date_time_string(), fmt % args))
|
|
|
|
def _send(
|
|
self,
|
|
status: int,
|
|
body: bytes,
|
|
content_type: str = "application/json; charset=utf-8",
|
|
headers: dict[str, str] | None = None,
|
|
) -> None:
|
|
self.send_response(int(status))
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Content-Length", str(len(body)))
|
|
for key, value in (headers or {}).items():
|
|
self.send_header(key, value)
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def _send_json(self, payload: Any, status: int = HTTPStatus.OK) -> None:
|
|
self._send(int(status), _json_bytes(payload))
|
|
|
|
def _send_error_json(self, status: int, message: str) -> None:
|
|
self._send_json({"error": message}, status)
|
|
|
|
def do_HEAD(self) -> None:
|
|
parsed = urlparse(self.path)
|
|
target = STATIC_DIR / "index.html" if parsed.path == "/" else STATIC_DIR / parsed.path.removeprefix("/static/")
|
|
if parsed.path == "/" or parsed.path.startswith("/static/"):
|
|
target = target.resolve()
|
|
static_root = STATIC_DIR.resolve()
|
|
if static_root not in target.parents and target != static_root:
|
|
self.send_response(HTTPStatus.FORBIDDEN)
|
|
self.end_headers()
|
|
return
|
|
if target.exists() and target.is_file():
|
|
content_type = mimetypes.guess_type(target.name)[0] or "application/octet-stream"
|
|
if target.suffix == ".html":
|
|
content_type = "text/html; charset=utf-8"
|
|
self.send_response(HTTPStatus.OK)
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Content-Length", str(target.stat().st_size))
|
|
self.end_headers()
|
|
return
|
|
self.send_response(HTTPStatus.NOT_FOUND)
|
|
self.end_headers()
|
|
|
|
def _read_json(self) -> dict[str, Any]:
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
if length <= 0:
|
|
return {}
|
|
raw = self.rfile.read(length).decode("utf-8")
|
|
return json.loads(raw)
|
|
|
|
def do_GET(self) -> None:
|
|
parsed = urlparse(self.path)
|
|
path = parsed.path
|
|
query = parse_qs(parsed.query)
|
|
try:
|
|
if path == "/":
|
|
self._serve_static("index.html")
|
|
elif path.startswith("/static/"):
|
|
self._serve_static(path.removeprefix("/static/"))
|
|
elif path == "/asset":
|
|
rel = unquote(query.get("path", [""])[0])
|
|
self._serve_asset(rel)
|
|
elif path == "/api/images":
|
|
self._send_json({"images": pipeline.list_images()})
|
|
elif path == "/api/capabilities":
|
|
self._send_json(pipeline.capabilities())
|
|
elif path == "/api/results":
|
|
filename = query.get("image", [""])[0]
|
|
self._send_json(pipeline.get_results(filename))
|
|
elif path == "/api/download":
|
|
self._serve_download(query)
|
|
elif path == "/api/job":
|
|
job_id = query.get("id", [""])[0]
|
|
with JOB_LOCK:
|
|
job = JOBS.get(job_id)
|
|
payload = dict(job) if job else None
|
|
if payload is None:
|
|
self._send_error_json(HTTPStatus.NOT_FOUND, "job not found")
|
|
else:
|
|
self._send_json(payload)
|
|
else:
|
|
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
|
|
except Exception as exc:
|
|
self._send_error_json(HTTPStatus.INTERNAL_SERVER_ERROR, str(exc))
|
|
|
|
def do_POST(self) -> None:
|
|
parsed = urlparse(self.path)
|
|
try:
|
|
payload = self._read_json()
|
|
if parsed.path == "/api/run":
|
|
job = _new_job("dehaze", _run_dehaze_job, payload)
|
|
self._send_json({"job": job})
|
|
elif parsed.path == "/api/run-default":
|
|
job = _new_job("default_batch", _run_default_batch_job, payload)
|
|
self._send_json({"job": job})
|
|
elif parsed.path == "/api/postprocess":
|
|
job = _new_job("postprocess", _run_post_job, payload)
|
|
self._send_json({"job": job})
|
|
else:
|
|
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
|
|
except Exception as exc:
|
|
self._send_error_json(HTTPStatus.BAD_REQUEST, str(exc))
|
|
|
|
def _serve_static(self, name: str) -> None:
|
|
target = (STATIC_DIR / name).resolve()
|
|
if STATIC_DIR.resolve() not in target.parents and target != STATIC_DIR.resolve():
|
|
self._send_error_json(HTTPStatus.FORBIDDEN, "forbidden")
|
|
return
|
|
if not target.exists() or not target.is_file():
|
|
self._send_error_json(HTTPStatus.NOT_FOUND, "static file not found")
|
|
return
|
|
content_type = mimetypes.guess_type(target.name)[0] or "application/octet-stream"
|
|
if target.suffix == ".html":
|
|
content_type = "text/html; charset=utf-8"
|
|
elif target.suffix == ".css":
|
|
content_type = "text/css; charset=utf-8"
|
|
elif target.suffix == ".js":
|
|
content_type = "application/javascript; charset=utf-8"
|
|
self._send(HTTPStatus.OK, target.read_bytes(), content_type)
|
|
|
|
def _serve_asset(self, rel: str) -> None:
|
|
target = pipeline.resolve_asset(rel)
|
|
content_type = mimetypes.guess_type(target.name)[0] or "application/octet-stream"
|
|
self._send(HTTPStatus.OK, target.read_bytes(), content_type)
|
|
|
|
def _serve_download(self, query: dict[str, list[str]]) -> None:
|
|
scope = query.get("scope", ["all"])[0]
|
|
if scope == "current":
|
|
filename = query.get("image", [""])[0]
|
|
if not filename:
|
|
self._send_error_json(HTTPStatus.BAD_REQUEST, "missing image")
|
|
return
|
|
items = pipeline.collect_download_items([filename])
|
|
archive_name = f"dehaze_{pipeline.image_id(filename)}.zip"
|
|
else:
|
|
items = pipeline.collect_download_items()
|
|
archive_name = "dehaze_results_all.zip"
|
|
|
|
if not items:
|
|
self._send_error_json(HTTPStatus.NOT_FOUND, "no downloadable files")
|
|
return
|
|
|
|
buffer = io.BytesIO()
|
|
with zipfile.ZipFile(buffer, "w", compression=zipfile.ZIP_DEFLATED) as archive:
|
|
for path, inner_name in items:
|
|
archive.write(path, inner_name)
|
|
|
|
self._send(
|
|
HTTPStatus.OK,
|
|
buffer.getvalue(),
|
|
"application/zip",
|
|
headers={"Content-Disposition": f'attachment; filename="{archive_name}"'},
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Local dehaze web console")
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, default=7860)
|
|
args = parser.parse_args()
|
|
|
|
server = ThreadingHTTPServer((args.host, args.port), DehazeRequestHandler)
|
|
print(f"Dehaze web console: http://{args.host}:{args.port}")
|
|
print(f"Images: {pipeline.IMAGE_DIR}")
|
|
print(f"Results: {pipeline.RESULTS_DIR}")
|
|
server.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|