src/perception/pipeline.py

"""感知层入口：图像审核 + OCR 并行"""
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from PIL import Image

from ..preprocess.scene_classifier import SceneResult
from .nsfw_detector import NSFWResult, detect as nsfw_detect
from .qr_detector import QRResult, detect as qr_detect
from .ocr_extractor import OCRResult, extract as ocr_extract


@dataclass
class PerceptionResult:
    nsfw: NSFWResult
    qr: QRResult
    ocr: OCRResult


def run(img: Image.Image, scene_result: SceneResult, debug_prefix: str = None) -> PerceptionResult:
    # print(f'scene_result {scene_result}')
    ocr = ocr_extract(img, scene_result, debug_prefix)
    # print(f"OCR extracted {len(ocr.blocks)} ocr {ocr}")
    with ThreadPoolExecutor(max_workers=2) as executor:
        f_nsfw = executor.submit(nsfw_detect, img)
        f_qr = executor.submit(qr_detect, img)
    # f_nsfw = nsfw_detect(img)
    # f_qr = qr_detect(img)
    nsfw = f_nsfw.result()
    qr = f_qr.result()
    if ocr.figure_images:
        with ThreadPoolExecutor(max_workers=4) as executor:
            figure_nsfw = list(executor.map(nsfw_detect, ocr.figure_images))
        max_figure = max(figure_nsfw, key=lambda r: r.score)
        if max_figure.score > nsfw.score:
            nsfw = max_figure

    return PerceptionResult(nsfw=nsfw, qr=qr, ocr=ocr)