import os from pathlib import Path from uuid import uuid4 from flask import Flask, render_template, request, send_from_directory, url_for from PIL import Image, ImageOps, UnidentifiedImageError from werkzeug.utils import secure_filename from src.decision.pipeline import decide from src.perception.pipeline import run as perception_run from src.preprocess.pipeline import run as preprocess_run from src.understanding.pipeline import run as understanding_run BASE_DIR = Path(__file__).resolve().parent IMAGES_DIR = BASE_DIR / "images" UPLOADS_DIR = BASE_DIR / "uploads" ALLOWED_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} UPLOADS_DIR.mkdir(exist_ok=True) app = Flask(__name__) app.config["MAX_CONTENT_LENGTH"] = 20 * 1024 * 1024 def _allowed_file(filename: str) -> bool: return Path(filename).suffix.lower() in ALLOWED_SUFFIXES def _default_image_name() -> str: default_name = "test.jpg" if (IMAGES_DIR / default_name).exists(): return default_name for image_path in sorted(IMAGES_DIR.iterdir()): if image_path.is_file() and _allowed_file(image_path.name): return image_path.name raise FileNotFoundError("latest/images/ 下没有可用的测试图片。") def _sample_image_names() -> list[str]: samples = [ image_path.name for image_path in sorted(IMAGES_DIR.iterdir()) if image_path.is_file() and _allowed_file(image_path.name) ] if not samples: raise FileNotFoundError("latest/images/ 下没有可用的测试图片。") return samples def _guess_format(filename: str, image: Image.Image) -> str: if image.format: return image.format.upper() suffix = Path(filename).suffix.lower() if suffix == ".png": return "PNG" if suffix in {".jpg", ".jpeg"}: return "JPEG" if suffix == ".webp": return "WEBP" if suffix == ".bmp": return "BMP" return "JPEG" def _decision_priority(action: str) -> int: return {"REJECT": 0, "REVIEW": 1, "PASS": 2}[action] def _enum_value(value): return getattr(value, "value", value) def _serialize_decision(result) -> dict: return { "action": _enum_value(result.action), "risk_level": result.risk_level, "primary_reason": result.primary_reason, "labels": list(result.labels), "score": round(float(result.score), 4), "evidence": list(result.evidence), } def _serialize_rule_hits(rule_hits: list) -> list[dict]: return [ { "rule_id": hit.rule_id, "level": _enum_value(hit.level), "matched_text": hit.matched_text, } for hit in rule_hits ] def _run_analysis(image_path: Path, source_name: str, preview_url: str, display_name: str | None = None) -> dict: with Image.open(image_path) as opened: fmt = _guess_format(source_name, opened) img = ImageOps.exif_transpose(opened).convert("RGB") pre = preprocess_run(img, fmt) slice_results = [] decisions = [] for index, sub_img in enumerate(pre.images, start=1): perc = perception_run(sub_img, pre.scene_result) texts = [block.text for block in perc.ocr.blocks] und = understanding_run(texts) decision = decide( scene=pre.scene_result.scene, nsfw=perc.nsfw, qr=perc.qr, ocr=perc.ocr, rule_hits=und.rule_hits, ) decisions.append(decision) slice_results.append( { "index": index, "decision": _serialize_decision(decision), "ocr_texts": texts, "ocr_block_count": len(perc.ocr.blocks), "ocr_avg_score": round(float(perc.ocr.avg_score), 4), "ocr_low_confidence": perc.ocr.low_confidence, "ocr_text_state": _enum_value(perc.ocr.text_state), "nsfw_score": round(float(perc.nsfw.score), 4), "nsfw_is_nsfw": perc.nsfw.is_nsfw, "qr_decision": _enum_value(perc.qr.decision), "qr_domains": list(perc.qr.domains), "qr_raw_contents": list(perc.qr.raw_contents), "normalized_text": und.normalized_text, "rule_hits": _serialize_rule_hits(und.rule_hits), } ) final_decision = min(decisions, key=lambda item: _decision_priority(_enum_value(item.action))) return { "source_name": display_name or source_name, "format": fmt, "image_size": {"width": img.width, "height": img.height}, "preview_url": preview_url, "scene": _enum_value(pre.scene_result.scene), "ocr_threshold": pre.scene_result.ocr_threshold, "warnings": list(pre.warnings), "slice_count": len(slice_results), "decision": _serialize_decision(final_decision), "slices": slice_results, } @app.route("/files//") def serve_file(kind: str, filename: str): directories = { "sample": IMAGES_DIR, "upload": UPLOADS_DIR, } directory = directories.get(kind) if directory is None: return ("Not Found", 404) return send_from_directory(directory, filename) @app.route("/default-image") def serve_default_image(): default_image = _default_image_name() return send_from_directory(IMAGES_DIR, default_image) @app.route("/", methods=["GET", "POST"]) def index(): error = None analysis = None default_image = _default_image_name() sample_images = _sample_image_names() selected_sample = default_image try: if request.method == "POST": analysis = None upload = request.files.get("image_file") selected_sample = request.form.get("sample_name", default_image) if selected_sample not in sample_images: raise ValueError("内置样例不存在,请重新选择。") if upload and upload.filename: if not _allowed_file(upload.filename): raise ValueError("仅支持 jpg/jpeg/png/bmp/webp 图片。") safe_name = secure_filename(upload.filename) or "upload.jpg" stored_name = f"{uuid4().hex[:8]}-{safe_name}" save_path = UPLOADS_DIR / stored_name upload.save(save_path) preview_url = url_for("serve_file", kind="upload", filename=stored_name) analysis = _run_analysis(save_path, upload.filename, preview_url) else: sample_path = IMAGES_DIR / selected_sample preview_url = url_for("serve_file", kind="sample", filename=selected_sample) analysis = _run_analysis(sample_path, selected_sample, preview_url, display_name=f"内置样例:{selected_sample}") else: sample_path = IMAGES_DIR / default_image preview_url = url_for("serve_default_image") analysis = _run_analysis(sample_path, default_image, preview_url, display_name="默认测试图") except (ValueError, UnidentifiedImageError, FileNotFoundError) as exc: error = str(exc) return render_template( "index.html", analysis=analysis, error=error, sample_images=sample_images, selected_sample=selected_sample, ) if __name__ == "__main__": host = os.getenv("HOST", "0.0.0.0") port = int(os.getenv("PORT", "5000")) app.run(host=host, port=port, debug=False)