safevision_utils.py

import os
import platform
import shutil
import tempfile

import cv2
import numpy as np


DEFAULT_CONTENT_LABELS = [
    "FEMALE_GENITALIA_COVERED",
    "FACE_FEMALE",
    "BUTTOCKS_EXPOSED",
    "FEMALE_BREAST_EXPOSED",
    "FEMALE_GENITALIA_EXPOSED",
    "MALE_BREAST_EXPOSED",
    "ANUS_EXPOSED",
    "FEET_EXPOSED",
    "BELLY_COVERED",
    "FEET_COVERED",
    "ARMPITS_COVERED",
    "ARMPITS_EXPOSED",
    "FACE_MALE",
    "BELLY_EXPOSED",
    "MALE_GENITALIA_EXPOSED",
    "ANUS_COVERED",
    "FEMALE_BREAST_COVERED",
    "BUTTOCKS_COVERED",
]

SAFETY_OBJECT_LABELS = [
    "cigarette",
    "cigar",
    "vape",
    "smoking_pipe",
    "joint",
    "alcohol_bottle",
    "beer_bottle",
    "wine_glass",
    "beer_glass",
    "cocktail_glass",
    "pill",
    "pill_bottle",
    "syringe",
    "cannabis_leaf",
    "drug_bag",
]

ALL_CENSOR_LABELS = DEFAULT_CONTENT_LABELS + SAFETY_OBJECT_LABELS


def label_group(label):
    label = str(label or "").upper()
    if label.startswith("FACE_"):
        return "face"
    if "COVERED" in label:
        return "covered"
    if "EXPOSED" in label:
        return "exposed"
    return "other"


def label_matches_filter(label, label_filter="exposed"):
    label_filter = str(label_filter or "exposed").lower()
    group = label_group(label)
    if label_filter == "all":
        return True
    if label_filter == "body":
        return group != "face"
    return group == "exposed"


def default_blur_rules(labels=None, blur=True):
    return {label: bool(blur) for label in (labels or ALL_CENSOR_LABELS)}


def detection_is_censorable(detection):
    if isinstance(detection, dict):
        if "censor" in detection:
            return bool(detection.get("censor"))
        label = detection.get("class", "")
        category = str(detection.get("category", "")).lower()
        if category in {"smoking", "alcohol", "drugs"}:
            return True
    else:
        label = detection

    label_text = str(label or "")
    return "EXPOSED" in label_text.upper() or label_text in SAFETY_OBJECT_LABELS


def parse_detector_selection(value=None, default="nude"):
    raw_value = value
    if raw_value in (None, ""):
        raw_value = os.environ.get("SAFEVISION_DETECTORS", default)

    selected = []
    for token in str(raw_value or default).replace(";", ",").replace("+", ",").split(","):
        name = token.strip().lower()
        if not name or name == "none":
            continue
        if name in {"all", "both", "combined"}:
            selected.extend(["nude", "objects"])
        elif name in {"nude", "nudity", "body", "safevision"}:
            selected.append("nude")
        elif name in {"object", "objects", "safety", "safety_objects", "cigarette", "smoking", "alcohol", "drugs"}:
            selected.append("objects")

    selected = list(dict.fromkeys(selected))
    if not selected:
        selected = list(dict.fromkeys(parse_detector_selection(default, default="nude")))
    return selected


def write_blur_exception_rules(path="BlurException.rule", rules=None, labels=None):
    path = os.fspath(path or "BlurException.rule")
    folder = os.path.dirname(os.path.abspath(path))
    if folder:
        os.makedirs(folder, exist_ok=True)

    rules = rules or default_blur_rules(labels)
    with open(path, "w", encoding="utf-8") as rule_file:
        for label in (labels or ALL_CENSOR_LABELS):
            rule_file.write(f"{label} = {'true' if rules.get(label, True) else 'false'}\n")
    return path


def ensure_blur_exception_rules(path="BlurException.rule", labels=None):
    path = os.fspath(path or "BlurException.rule")
    if not os.path.exists(path):
        write_blur_exception_rules(path, labels=labels)
        print(f"Created default blur exception rules at: {path}")
    return path


def load_blur_exception_rules(path="BlurException.rule", labels=None):
    path = ensure_blur_exception_rules(path, labels=labels)
    rules = default_blur_rules(labels)
    with open(path, "r", encoding="utf-8") as rule_file:
        for line in rule_file:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            label, blur = line.split("=", 1)
            rules[label.strip()] = blur.strip().lower() in {"1", "true", "yes", "on"}
    return rules


def cv2_imread(path, flags=cv2.IMREAD_COLOR):
    """Read an image through imdecode so Windows Unicode paths work."""
    path = os.fspath(path)
    try:
        data = np.fromfile(path, dtype=np.uint8)
        if data.size == 0:
            return None
        return cv2.imdecode(data, flags)
    except Exception:
        return cv2.imread(path, flags)


def cv2_imwrite(path, image, params=None):
    """Write an image through imencode so Windows Unicode paths work."""
    path = os.fspath(path)
    folder = os.path.dirname(os.path.abspath(path))
    if folder:
        os.makedirs(folder, exist_ok=True)

    ext = os.path.splitext(path)[1] or ".jpg"
    try:
        ok, encoded = cv2.imencode(ext, image, params or [])
        if not ok:
            return False
        encoded.tofile(path)
        return True
    except Exception:
        return cv2.imwrite(path, image, params or [])


def normalize_mask_shape(mask_shape="rectangle"):
    value = str(mask_shape or "rectangle").strip().lower()
    if value in {"ellipse", "oval", "circle", "round"}:
        return "ellipse"
    return "rectangle"


def make_blur_kernel(strength=None, sigma=None, default=(23, 23, 30)):
    if strength in (None, ""):
        return default

    try:
        kernel = int(float(strength))
    except (TypeError, ValueError):
        return default

    kernel = max(3, min(kernel, 151))
    if kernel % 2 == 0:
        kernel += 1

    if sigma in (None, ""):
        sigma_value = max(1.0, float(kernel))
    else:
        try:
            sigma_value = max(1.0, float(sigma))
        except (TypeError, ValueError):
            sigma_value = max(1.0, float(kernel))

    return (kernel, kernel, sigma_value)


def apply_region_censor(
    image,
    x,
    y,
    w,
    h,
    blur_kernel=(23, 23, 30),
    use_solid_color=False,
    solid_color=(0, 0, 0),
    mask_shape="rectangle",
):
    image_height, image_width = image.shape[:2]
    x1 = max(0, int(x))
    y1 = max(0, int(y))
    x2 = min(image_width, int(x + w))
    y2 = min(image_height, int(y + h))
    if x2 <= x1 or y2 <= y1:
        return False

    roi = image[y1:y2, x1:x2]
    roi_height, roi_width = roi.shape[:2]

    if use_solid_color:
        censored_roi = np.full((roi_height, roi_width, 3), solid_color, dtype=np.uint8)
    else:
        kernel_x, kernel_y, kernel_sigma = blur_kernel
        censored_roi = cv2.GaussianBlur(roi, (int(kernel_x), int(kernel_y)), float(kernel_sigma))

    if normalize_mask_shape(mask_shape) == "ellipse":
        mask = np.zeros((roi_height, roi_width), dtype=np.uint8)
        center = (roi_width // 2, roi_height // 2)
        axes = (max(1, roi_width // 2), max(1, roi_height // 2))
        cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
        roi[mask > 0] = censored_roi[mask > 0]
    else:
        image[y1:y2, x1:x2] = censored_roi

    return True


class ManagedVideoCapture:
    """cv2.VideoCapture wrapper with a Windows Unicode-path fallback."""

    def __init__(self, path):
        self.original_path = os.fspath(path)
        self.path_in_use = self.original_path
        self._temp_dir = None
        self._capture = cv2.VideoCapture(self.original_path)

        if not self._capture.isOpened() and platform.system() == "Windows" and os.path.exists(self.original_path):
            self._capture.release()
            suffix = os.path.splitext(self.original_path)[1] or ".mp4"
            self._temp_dir = tempfile.TemporaryDirectory(prefix="safevision_video_")
            self.path_in_use = os.path.join(self._temp_dir.name, f"input{suffix}")
            shutil.copy2(self.original_path, self.path_in_use)
            self._capture = cv2.VideoCapture(self.path_in_use)

    def __getattr__(self, name):
        return getattr(self._capture, name)

    def release(self):
        self._capture.release()
        if self._temp_dir is not None:
            self._temp_dir.cleanup()
            self._temp_dir = None


def open_video_capture(path):
    return ManagedVideoCapture(path)


def parse_provider_list(value):
    if not value:
        env_value = os.environ.get("SAFEVISION_ONNX_PROVIDERS", "")
        value = env_value
    if not value:
        return None
    return [provider.strip() for provider in value.split(",") if provider.strip()]


def select_onnx_providers(requested=None):
    import onnxruntime

    available = onnxruntime.get_available_providers()
    if requested:
        selected = [provider for provider in requested if provider in available]
        missing = [provider for provider in requested if provider not in available]
        if missing:
            print(f"Requested ONNX providers are not available and will be skipped: {missing}")
        if "CPUExecutionProvider" in available and "CPUExecutionProvider" not in selected:
            selected.append("CPUExecutionProvider")
        return selected or ["CPUExecutionProvider"]

    allow_tensorrt = os.environ.get("SAFEVISION_ENABLE_TENSORRT", "").lower() in {"1", "true", "yes"}
    preferred = []
    if allow_tensorrt:
        preferred.append("TensorrtExecutionProvider")
    preferred.extend([
        "CUDAExecutionProvider",
        "DmlExecutionProvider",
        "DirectMLExecutionProvider",
        "OpenVINOExecutionProvider",
        "ROCMExecutionProvider",
        "CoreMLExecutionProvider",
        "CPUExecutionProvider",
    ])

    selected = []
    for provider in preferred:
        if provider in available and provider not in selected:
            selected.append(provider)

    if not selected:
        selected = available or ["CPUExecutionProvider"]

    skipped_tensorrt = "TensorrtExecutionProvider" in available and not allow_tensorrt
    if skipped_tensorrt:
        print("TensorRT provider detected but disabled by default. Set SAFEVISION_ENABLE_TENSORRT=1 to opt in.")
    print(f"Using ONNX Runtime providers: {selected}")
    return selected


def create_onnx_session(model_path, providers=None, sess_options=None):
    import onnxruntime

    selected = select_onnx_providers(providers)
    try:
        if sess_options is not None:
            return onnxruntime.InferenceSession(model_path, sess_options=sess_options, providers=selected)
        return onnxruntime.InferenceSession(model_path, providers=selected)
    except Exception as exc:
        if selected != ["CPUExecutionProvider"] and "CPUExecutionProvider" in onnxruntime.get_available_providers():
            print(f"ONNX provider initialization failed ({exc}). Falling back to CPUExecutionProvider.")
            if sess_options is not None:
                return onnxruntime.InferenceSession(model_path, sess_options=sess_options, providers=["CPUExecutionProvider"])
            return onnxruntime.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        raise