app.py

# app.py - SMART RULE-BASED COMBINE SYSTEM
# Your 7 custom rules for maximum accuracy!

from flask import Flask, request, jsonify
from transformers import (
    AutoModelForImageClassification,
    ViTImageProcessor
)
from PIL import Image
import torch
import io
import base64
import numpy as np
import os
import time
import json
import shutil
import concurrent.futures

app = Flask(__name__)

# ============================================================
# MODEL LOADING
# ============================================================
print("=" * 55)

# Falconsai ViT
print("Loading Falconsai ViT...")
FALCON_OK = False
falcon_model = None
falcon_processor = None
try:
    falcon_processor = ViTImageProcessor.from_pretrained(
        "Falconsai/nsfw_image_detection")
    falcon_model = AutoModelForImageClassification.from_pretrained(
        "Falconsai/nsfw_image_detection")
    falcon_model.eval()
    FALCON_OK = True
    print("✅ Falconsai ViT OK!")
except Exception as e:
    print(f"❌ Falconsai ViT FAILED: {e}")

# AdamCodd ViT
print("Loading AdamCodd ViT...")
ADAM_OK = False
adam_model = None
adam_processor = None
try:
    adam_processor = ViTImageProcessor.from_pretrained(
        "AdamCodd/vit-base-nsfw-detector")
    adam_model = AutoModelForImageClassification.from_pretrained(
        "AdamCodd/vit-base-nsfw-detector")
    adam_model.eval()
    ADAM_OK = True
    print("✅ AdamCodd ViT OK!")
except Exception as e:
    print(f"❌ AdamCodd ViT FAILED: {e}")

# EraX V1.1 Medium
print("Loading EraX V1.1 MEDIUM...")
erax_v11_model = None
ERAX_V11_OK = False
try:
    from huggingface_hub import hf_hub_download
    from ultralytics import YOLO
    os.makedirs("./models", exist_ok=True)
    path_v11 = hf_hub_download(
        repo_id="erax-ai/EraX-Anti-NSFW-V1.1",
        filename="erax-anti-nsfw-yolo11m-v1.1.pt",
        local_dir="./models")
    erax_v11_model = YOLO(path_v11)
    ERAX_V11_OK = True
    print("✅ EraX V1.1 MEDIUM OK!")
except Exception as e:
    print(f"❌ EraX V1.1 FAILED: {e}")
    try:
        path_v11n = hf_hub_download(
            repo_id="erax-ai/EraX-Anti-NSFW-V1.1",
            filename="erax-anti-nsfw-yolo11n-v1.1.pt",
            local_dir="./models")
        erax_v11_model = YOLO(path_v11n)
        ERAX_V11_OK = True
        print("✅ EraX V1.1 NANO OK (fallback)")
    except:
        pass

# EraX V1.0 Medium
print("Loading EraX V1.0 MEDIUM...")
erax_v10_model = None
ERAX_V10_OK = False
try:
    path_v10 = hf_hub_download(
        repo_id="erax-ai/EraX-NSFW-V1.0",
        filename="erax_nsfw_yolo11m.pt",
        local_dir="./models")
    erax_v10_model = YOLO(path_v10)
    ERAX_V10_OK = True
    print("✅ EraX V1.0 MEDIUM OK!")
except Exception as e:
    print(f"❌ EraX V1.0 FAILED: {e}")

# Falconsai YOLOv9 ONNX
print("Loading Falconsai YOLOv9 ONNX...")
falcon_yolo_session = None
falcon_yolo_labels = None
FALCON_YOLO_OK = False
try:
    import onnxruntime as ort
    pt_path = hf_hub_download(
        repo_id="Falconsai/nsfw_image_detection",
        filename="falconsai_yolov9_nsfw_model_quantized.pt",
        local_dir="./models")
    onnx_path = pt_path.replace('.pt', '.onnx')
    shutil.copy2(pt_path, onnx_path)
    falcon_yolo_session = ort.InferenceSession(
        onnx_path, providers=['CPUExecutionProvider'])
    labels_path = hf_hub_download(
        repo_id="Falconsai/nsfw_image_detection",
        filename="labels.json", local_dir="./models")
    with open(labels_path) as f:
        falcon_yolo_labels = json.load(f)
    print(f"✅ Falconsai YOLOv9 OK! Labels: {falcon_yolo_labels}")
    FALCON_YOLO_OK = True
except Exception as e:
    print(f"❌ Falconsai YOLOv9 FAILED: {e}")

print("=" * 55)
print(f"Falconsai ViT  : {'✅' if FALCON_OK     else '❌'}")
print(f"AdamCodd ViT   : {'✅' if ADAM_OK       else '❌'}")
print(f"EraX V1.1 Med  : {'✅' if ERAX_V11_OK   else '❌'}")
print(f"EraX V1.0 Med  : {'✅' if ERAX_V10_OK   else '❌'}")
print(f"Falcon YOLOv9  : {'✅' if FALCON_YOLO_OK else '❌'}")
print("🛡️ Server ready!")
print("=" * 55)

ERAX_CLASSES = {
    0: "anus",
    1: "make_love",
    2: "nipple",
    3: "penis",
    4: "vagina"
}

# Vulgar words for EraX detection
VULGAR_WORDS = {"anus", "make_love", "penis", "vagina"}
ALL_WORDS    = {"anus", "make_love", "nipple", "penis", "vagina"}

# ============================================================
# PREPROCESSING
# ============================================================

def to_square(img):
    w, h = img.size
    s    = min(w, h)
    return img.crop(((w-s)//2, (h-s)//2,
                     (w-s)//2+s, (h-s)//2+s))

def to_size(img, size):
    return img.resize((size, size), Image.LANCZOS)

def get_tiles(image, model_size):
    """
    3 tiles — all run in PARALLEL!
    Tile 1: Full image → square → resize
    Tile 2: Center 60% → square → resize
    Tile 3: Center 40% → square → resize
    """
    w, h  = image.size
    tiles = []

    # Tile 1 — full image
    tiles.append(("full",
        to_size(to_square(image), model_size)))

    # Tile 2 — center 60%
    m  = 0.20
    c2 = image.crop((int(w*m), int(h*m),
                     int(w*(1-m)), int(h*(1-m))))
    if c2.width > 80:
        tiles.append(("center_60",
            to_size(to_square(c2), model_size)))

    # Tile 3 — center 40%
    m2 = 0.30
    c3 = image.crop((int(w*m2), int(h*m2),
                     int(w*(1-m2)), int(h*(1-m2))))
    if c3.width > 60:
        tiles.append(("center_40",
            to_size(to_square(c3), model_size)))

    return tiles

# ============================================================
# INFERENCE
# ============================================================

def score_vit(tile, model, processor):
    inputs = processor(images=tile, return_tensors="pt")
    with torch.no_grad():
        out = model(**inputs)
    probs    = torch.softmax(out.logits, dim=1)[0]
    id2label = model.config.id2label
    scores   = {id2label[i].lower(): float(p)
                for i, p in enumerate(probs)}
    nsfw = scores.get('nsfw',
           scores.get('unsafe',
           1.0 - scores.get('normal',
           1.0 - scores.get('sfw', 1.0))))
    return round(nsfw, 4)

def run_vit_parallel(image, model, processor,
                     model_size, name):
    """
    Run ALL 3 tiles in PARALLEL using threads.
    Total time = slowest tile (not sum of 3).
    Best score across all 3 tiles returned.
    """
    if model is None or processor is None:
        return 0.0
    try:
        t0    = time.time()
        tiles = get_tiles(image, model_size)

        def score_one(args):
            tile_name, tile = args
            s = score_vit(tile, model, processor)
            print(f"    {name}[{tile_name}]: {s:.1%}")
            return s

        # Run all 3 tiles simultaneously!
        with concurrent.futures.ThreadPoolExecutor(
                max_workers=3) as ex:
            scores = list(ex.map(
                score_one, tiles,
                timeout=15))

        best = round(max(scores), 4)
        print(f"  {name} BEST: {best:.1%} "
              f"({time.time()-t0:.2f}s) ⚡parallel")
        return best

    except Exception as e:
        print(f"{name} parallel error: {e}")
        return 0.0

def run_falconsai(image):
    """Falconsai ViT — 3 tiles in parallel"""
    if not FALCON_OK or falcon_model is None:
        return 0.0
    return run_vit_parallel(
        image, falcon_model,
        falcon_processor, 224, "FalconViT")

def run_adamcodd(image):
    """AdamCodd ViT — 3 tiles in parallel"""
    if not ADAM_OK or adam_model is None:
        return 0.0
    return run_vit_parallel(
        image, adam_model,
        adam_processor, 384, "AdamCodd")

def get_image_tiles_3(image):
    """3 tiles: full + center60 + center40"""
    w, h  = image.size
    tiles = []
    tiles.append(("full", image))
    m  = 0.20
    c2 = image.crop((int(w*m), int(h*m),
                     int(w*(1-m)), int(h*(1-m))))
    if c2.width > 80 and c2.height > 80:
        tiles.append(("center_60", c2))
    m2 = 0.30
    c3 = image.crop((int(w*m2), int(h*m2),
                     int(w*(1-m2)), int(h*(1-m2))))
    if c3.width > 60 and c3.height > 60:
        tiles.append(("center_40", c3))
    return tiles

def get_image_tiles_2(image):
    """2 tiles: full + center60 only"""
    w, h  = image.size
    tiles = []
    tiles.append(("full", image))
    m  = 0.20
    c2 = image.crop((int(w*m), int(h*m),
                     int(w*(1-m)), int(h*(1-m))))
    if c2.width > 80 and c2.height > 80:
        tiles.append(("center_60", c2))
    return tiles

def run_erax_one_tile(tile_img, model):
    """Run EraX on one image tile"""
    img_array = np.array(tile_img.convert("RGB"))
    results   = model(img_array, conf=0.20,
                      iou=0.30, verbose=False)
    dets   = []
    counts = {}
    for r in results:
        if r.boxes is not None:
            for box in r.boxes:
                cid  = int(box.cls[0])
                conf = round(float(box.conf[0]), 4)
                cls  = ERAX_CLASSES.get(cid, "unknown")
                dets.append({"class": cls, "confidence": conf})
                counts[cls] = counts.get(cls, 0) + 1
    return dets, counts

def run_erax(image, model, name, n_tiles=2):
    """
    EraX YOLO — tiles in PARALLEL.
    n_tiles=2: EraX V1.1 (full + center60)
    n_tiles=0: EraX V1.0 (full image only — no tiles)
    Merges all detections from all tiles.
    """
    if model is None:
        return [], {}
    try:
        t0    = time.time()
        if n_tiles == 0:
            tiles = [("full", image)]
        elif n_tiles == 2:
            tiles = get_image_tiles_2(image)
        else:
            tiles = get_image_tiles_3(image)

        def process_tile(args):
            tile_name, tile_img = args
            dets, cnts = run_erax_one_tile(tile_img, model)
            if dets:
                print(f"    {name}[{tile_name}]: "
                      f"{[d['class'] for d in dets]}")
            return dets, cnts

        # Run tiles in PARALLEL
        max_w = min(len(tiles), 3)
        with concurrent.futures.ThreadPoolExecutor(
                max_workers=max_w) as ex:
            tile_results = list(ex.map(
                process_tile, tiles, timeout=15))

        # Merge all detections from all tiles
        all_dets   = []
        all_counts = {}
        seen       = set()  # deduplicate same class+conf

        for dets, cnts in tile_results:
            for d in dets:
                key = f"{d['class']}_{d['confidence']}"
                if key not in seen:
                    seen.add(key)
                    all_dets.append(d)
            for cls, cnt in cnts.items():
                all_counts[cls] = all_counts.get(cls, 0) + cnt

        print(f"  {name} MERGED: {all_dets} "
              f"({time.time()-t0:.2f}s) ⚡parallel")
        return all_dets, all_counts

    except Exception as e:
        print(f"{name} parallel error: {e}")
        return [], {}

def score_falcon_yolo9_one(tile_img):
    """Score one tile with FalconYOLO9 ONNX"""
    inp     = falcon_yolo_session.get_inputs()[0]
    inp_shp = inp.shape
    if len(inp_shp) == 4:
        _, _, h, w = inp_shp
        size = (int(w) if isinstance(w, int) and w > 0 else 640,
                int(h) if isinstance(h, int) and h > 0 else 640)
    else:
        size = (640, 640)

    img  = tile_img.convert("RGB").resize(size, Image.LANCZOS)
    arr  = np.array(img).astype(np.float32) / 255.0
    arr  = arr.transpose(2, 0, 1)
    arr  = np.expand_dims(arr, 0)
    outs = falcon_yolo_session.run(None, {inp.name: arr})
    out  = outs[0]

    max_nsfw = 0.0
    if out.ndim == 2:
        probs = out[0]
        exp   = np.exp(probs - probs.max())
        probs = exp / exp.sum()
        for i, p in enumerate(probs):
            lbl = str(falcon_yolo_labels.get(
                str(i), "")).lower()
            if 'nsfw' in lbl and float(p) > max_nsfw:
                max_nsfw = float(p)
    elif out.ndim == 3:
        conf = out[0, :, 4] if out.shape[2] > 4 \
               else out[0, :, 0]
        mx   = float(conf.max()) if len(conf) > 0 else 0.0
        if mx >= 0.20:
            max_nsfw = mx

    return round(max_nsfw, 4)

def run_falcon_yolo9(image):
    """
    FalconYOLO9 ONNX — 2 tiles in PARALLEL.
    (full + center60 only — faster)
    Best score across tiles returned.
    """
    if not FALCON_YOLO_OK or falcon_yolo_session is None:
        return 0.0
    try:
        t0    = time.time()
        tiles = get_image_tiles_2(image)

        def score_tile(args):
            tile_name, tile_img = args
            s = score_falcon_yolo9_one(tile_img)
            print(f"    FalconYOLO9[{tile_name}]: {s:.1%}")
            return s

        # Run 2 tiles in PARALLEL
        with concurrent.futures.ThreadPoolExecutor(
                max_workers=2) as ex:
            scores = list(ex.map(
                score_tile, tiles, timeout=15))

        best = round(max(scores), 4)
        print(f"  FalconYOLO9 BEST: {best:.1%} "
              f"({time.time()-t0:.2f}s) ⚡2-tile parallel")
        return best

    except Exception as e:
        print(f"FalconYOLO9 parallel error: {e}")
        return 0.0

# ============================================================
# YOUR 7 SMART RULES
# ============================================================

def combine_detections(e11_dets, e10_dets):
    """Merge detections from both EraX models"""
    all_dets   = e11_dets + e10_dets
    classes    = [d['class'] for d in all_dets]
    class_set  = set(classes)
    # Count occurrences of each word across both models
    counts = {}
    for c in classes:
        counts[c] = counts.get(c, 0) + 1
    return all_dets, class_set, counts

def apply_rules(fs, as_, fys,
                e11_dets, e11_counts,
                e10_dets, e10_counts):
    """
    Apply all 7 smart rules.
    Returns: (is_nsfw, confidence, rule_triggered, reason)
    """

    # Merge both EraX detections
    all_dets, class_set, all_counts = combine_detections(
        e11_dets, e10_dets)

    # Helper functions
    def has(word):
        return word in class_set

    def count(word):
        return all_counts.get(word, 0)

    def vulgar_detected():
        """Any vulgar word found (not nipple alone)"""
        return bool(class_set & VULGAR_WORDS)

    def any_erax():
        return bool(class_set & ALL_WORDS)

    # ─────────────────────────────────────────────────
    # RULE 1: Falconsai (90%+) + EraX vulgar word
    #
    # TRIGGER: Falconsai>=90% + ONE single vulgar word
    #          (anus / penis / vagina / make_love alone)
    #          BUT only if nipple is NOT also present
    #
    # SKIP: nipple alone           → false positive
    # SKIP: nipple + vagina        → false positive
    # SKIP: vagina + nipple        → false positive (same)
    # SKIP: any combo WITH nipple  → skip Rule 1
    #
    # NSFW: Falconsai>=90% + nipple × 8 or more
    # ─────────────────────────────────────────────────
    if fs >= 0.90:
        erax_vulgar = class_set & VULGAR_WORDS
        has_nipple  = has('nipple')
        nipple_count = count('nipple')

        if erax_vulgar and not has_nipple:
            # Single or multiple vulgar words found
            # BUT nipple is NOT present = NSFW
            return (True, fs, "Rule 1",
                f"Falconsai {fs:.0%} + EraX vulgar: {erax_vulgar}")

        elif erax_vulgar and has_nipple:
            # Vulgar word + nipple = SKIP
            # (nipple presence makes it ambiguous)
            print(f"  Rule 1 SKIPPED: vulgar+nipple combo")

        elif has_nipple and not erax_vulgar:
            # Nipple only (no other vulgar words)
            if nipple_count >= 8:
                # 8+ nipple detections = very confident NSFW
                return (True, fs, "Rule 1",
                    f"Falconsai {fs:.0%} + {nipple_count}x nipple (high count)")
            # nipple < 8 = SKIP
            print(f"  Rule 1 SKIPPED: nipple only ({nipple_count}x < 8)")

    # ─────────────────────────────────────────────────
    # RULE 2: AdamCodd (90%+) + FalconYOLO (50%+)
    # ─────────────────────────────────────────────────
    if as_ >= 0.90 and fys >= 0.50:
        return (True, max(as_, fys), "Rule 2",
            f"AdamCodd {as_:.0%} + FalconYOLO {fys:.0%}")

    # ─────────────────────────────────────────────────
    # RULE 3: Falconsai (90%+) + AdamCodd (90%+)
    #         + EraX (make_love or any vulgar)
    # ─────────────────────────────────────────────────
    if fs >= 0.90 and as_ >= 0.90:
        if has('make_love') or vulgar_detected():
            return (True, max(fs, as_), "Rule 3",
                f"Falconsai {fs:.0%} + AdamCodd {as_:.0%} "
                f"+ EraX: {class_set & (VULGAR_WORDS | {'make_love'})}")

    # ─────────────────────────────────────────────────
    # RULE 4: EraX body part combinations
    # make_love + vagina + nipple
    # vagina + anus
    # vagina + penis
    # anus + penis
    # 6+ vagina detections
    # ─────────────────────────────────────────────────
    if all_dets:
        vagina_count = count('vagina')

        # make_love + vagina + nipple
        if (has('make_love') and
            has('vagina') and has('nipple')):
            return (True, 0.95, "Rule 4a",
                "EraX: make_love + vagina + nipple")

        # vagina + anus
        if has('vagina') and has('anus'):
            return (True, 0.95, "Rule 4b",
                "EraX: vagina + anus")

        # vagina + penis
        if has('vagina') and has('penis'):
            return (True, 0.95, "Rule 4c",
                "EraX: vagina + penis")

        # anus + penis
        if has('anus') and has('penis'):
            return (True, 0.95, "Rule 4d",
                "EraX: anus + penis")

        # 6+ vagina detections
        if vagina_count >= 6:
            return (True, 0.95, "Rule 4e",
                f"EraX: {vagina_count}x vagina detected")

        # make_love alone REMOVED — not reliable enough

    # ─────────────────────────────────────────────────
    # RULE 5: AdamCodd (90%+) + EraX vulgar word
    # ─────────────────────────────────────────────────
    if as_ >= 0.90 and vulgar_detected():
        return (True, as_, "Rule 5",
            f"AdamCodd {as_:.0%} + EraX vulgar: "
            f"{class_set & VULGAR_WORDS}")

    # ─────────────────────────────────────────────────
    # RULE 6: Falconsai (90%+) + AdamCodd (90%+)
    #         + FalconYOLO (50%+)
    # ─────────────────────────────────────────────────
    if fs >= 0.90 and as_ >= 0.90 and fys >= 0.50:
        return (True, max(fs, as_, fys), "Rule 6",
            f"Falconsai {fs:.0%} + AdamCodd {as_:.0%} "
            f"+ FalconYOLO {fys:.0%}")

    # ─────────────────────────────────────────────────
    # RULE 7: Both ViT very high (93%+)
    # ─────────────────────────────────────────────────
    if fs >= 0.93 and as_ >= 0.93:
        return (True, max(fs, as_), "Rule 7",
            f"Falconsai {fs:.0%} + AdamCodd {as_:.0%} "
            f"both very high")

    # ─────────────────────────────────────────────────
    # RULE 8: Falconsai (90%+) + FalconYOLO (90%+)
    # ─────────────────────────────────────────────────
    if fs >= 0.90 and fys >= 0.90:
        return (True, max(fs, fys), "Rule 8",
            f"Falconsai {fs:.0%} + FalconYOLO {fys:.0%} both high")

    # No rule triggered
    return (False, max(fs, as_, fys), "None",
        "No rule triggered — SAFE")

# ============================================================
# PARALLEL MODEL EXECUTION
# ============================================================

def run_all_parallel(image):
    results = {}
    with concurrent.futures.ThreadPoolExecutor(
            max_workers=5) as executor:
        futures = {
            executor.submit(run_falconsai, image):               "fs",
            executor.submit(run_adamcodd,  image):               "as_",
            executor.submit(run_erax, image,
                erax_v11_model, "EraX-V1.1", 2):               "e11",
            executor.submit(run_erax, image,
                erax_v10_model, "EraX-V1.0", 0):               "e10",
            executor.submit(run_falcon_yolo9, image):            "fys"
        }
        for future in concurrent.futures.as_completed(
                futures, timeout=25):
            key = futures[future]
            try:
                results[key] = future.result(timeout=5)
            except Exception as e:
                print(f"  {key} parallel error: {e}")
                results[key] = ([], {}) if key in ["e11","e10"] else 0.0
    return results

# ============================================================
# ROUTES
# ============================================================

@app.route('/detect', methods=['POST'])
def detect():
    try:
        data = request.json
        if not data or 'image' not in data:
            return jsonify({'error': 'No image'}), 400

        img_bytes = base64.b64decode(data['image'])
        image     = Image.open(
            io.BytesIO(img_bytes)).convert('RGB')

        # ── Read source_url and tab_id sent by extension ──────
        # Bound to the screenshot at capture time.
        # Echoed back so extension blocks the correct URL
        # even if user switched tabs while server was processing.
        source_url = data.get('source_url', '')
        tab_id     = data.get('tab_id',     None)

        print(f"\n{'='*55}")
        print(f"Image     : {image.size[0]}x{image.size[1]}")
        print(f"Source URL: {source_url}")
        print(f"Tab ID    : {tab_id}")
        t0 = time.time()

        # Run all models in parallel
        res = run_all_parallel(image)

        fs   = res.get("fs",  0.0)
        as_  = res.get("as_", 0.0)
        fys  = res.get("fys", 0.0)
        e11_dets, e11_counts = res.get("e11", ([], {}))
        e10_dets, e10_counts = res.get("e10", ([], {}))

        # Apply your 7 smart rules
        (is_nsfw, confidence,
         rule, reason) = apply_rules(
            fs, as_, fys,
            e11_dets, e11_counts,
            e10_dets, e10_counts
        )

        t_total = round(time.time() - t0, 2)

        # Merged EraX info for response
        all_dets, class_set, all_counts = combine_detections(
            e11_dets, e10_dets)

        print(f"\n{'─'*55}")
        print(f"FalconViT  : {fs:.1%}")
        print(f"AdamCodd   : {as_:.1%}")
        print(f"EraX V1.1  : {[d['class'] for d in e11_dets]}")
        print(f"EraX V1.0  : {[d['class'] for d in e10_dets]}")
        print(f"FalconYOLO9: {fys:.1%}")
        print(f"Rule       : {rule}")
        print(f"Reason     : {reason}")
        print(f"Result     : {'🔴 NSFW' if is_nsfw else '🟢 SAFE'} "
              f"({confidence:.1%})")
        print(f"Time       : {t_total}s")
        print(f"{'='*55}\n")

        return jsonify({
            'nsfw':        is_nsfw,
            'confidence':  round(confidence, 3),
            'rule':        rule,
            'reason':      reason,
            'total_time':  t_total,

            # ── URL binding: echoed back to extension ──────────
            # Extension uses these to block the correct URL
            # regardless of which tab is active when this
            # response arrives.
            'source_url':  source_url,
            'tab_id':      tab_id,

            # Individual scores
            'falconsai_score':  fs,
            'adam_score':       as_,
            'falcon_yolo_score': fys,

            # EraX combined detections
            'erax_detections':  all_dets,
            'erax_classes':     list(class_set),
            'erax_counts':      all_counts,

            # Detailed per model
            'falconsai_vit':  {'score': fs, 'nsfw': fs >= 0.90},
            'adamcodd':       {'score': as_, 'nsfw': as_ >= 0.90},
            'erax_v11_medium': {
                'detections': e11_dets,
                'counts': e11_counts
            },
            'erax_v10_medium': {
                'detections': e10_dets,
                'counts': e10_counts
            },
            'falconsai_yolo9': {
                'score': fys, 'nsfw': fys >= 0.50
            },

            # Status
            'models_status': {
                'falcon_vit':  FALCON_OK,
                'adamcodd':    ADAM_OK,
                'erax_v11':    ERAX_V11_OK,
                'erax_v10':    ERAX_V10_OK,
                'falcon_yolo': FALCON_YOLO_OK
            }
        })

    except Exception as e:
        print(f"Detect error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/ping', methods=['GET'])
def ping():
    return jsonify({
        'status':      'alive',
        'falcon_vit':  FALCON_OK,
        'adamcodd':    ADAM_OK,
        'erax_v11':    ERAX_V11_OK,
        'erax_v10':    ERAX_V10_OK,
        'falcon_yolo': FALCON_YOLO_OK
    })

@app.route('/rules', methods=['GET'])
def rules():
    return jsonify({
        'rules': {
            'Rule 1': 'Falconsai>=90% + EraX pure vulgar word '
                      '(nipple alone = SKIP; nipple+vagina = SKIP; '
                      '6+nipple = NSFW)',
            'Rule 2': 'AdamCodd>=90% + FalconYOLO>=50%',
            'Rule 3': 'Falconsai>=90% + AdamCodd>=90% + EraX(make_love/vulgar)',
            'Rule 4': 'EraX combinations: '
                      'make_love+vagina+nipple / vagina+anus / vagina+penis / '
                      'anus+penis / 6+vagina '
                      '(make_love alone REMOVED)',
            'Rule 8': 'Falconsai>=90% + FalconYOLO>=90% both high',
            'Rule 5': 'AdamCodd>=90% + EraX vulgar word',
            'Rule 6': 'Falconsai>=90% + AdamCodd>=90% + FalconYOLO>=50%',
            'Rule 7': 'Falconsai>=93% + AdamCodd>=93% both very high'
        },
        'vulgar_words':  list(VULGAR_WORDS),
        'all_erax_classes': list(ALL_WORDS)
    })

@app.route('/', methods=['GET'])
def home():
    return jsonify({
        'status': '✅ Running — Smart 7-Rule System',
        'models': {
            'falconsai_vit':   '✅' if FALCON_OK     else '❌',
            'adamcodd_vit':    '✅' if ADAM_OK       else '❌',
            'erax_v11_medium': '✅' if ERAX_V11_OK   else '❌',
            'erax_v10_medium': '✅' if ERAX_V10_OK   else '❌',
            'falconsai_yolo9': '✅' if FALCON_YOLO_OK else '❌'
        },
        'docs': 'Visit /rules for all 7 detection rules'
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)