"""用 LocalAI qwen3.5-9b 过滤不含活体动物的图片。

用法：python filter_images.py
输入：datasets/images/<class>/*.jpg
输出：删除不含目标动物的图片
"""
import os
import json
import base64
import sys
import urllib.request
from io import BytesIO
from PIL import Image

# LocalAI 服务器配置
LOCAL_AI_URL = "http://192.168.3.210:18080/v1/chat/completions"
MODEL_NAME = "qwen3.5-9b"

# 过滤配置
BASE = os.path.dirname(__file__)
IMG_DIR = os.path.join(BASE, "datasets", "images")
CLASSES = ["pheasant"]
MAX_IMAGE_SIZE = 400  # 缩放到最大边长

# 过滤提示词
FILTER_PROMPT = """你是一个图片质量检查助手。请判断这张图片是否包含【{target}】的活体动物照片。

判断标准：
✓ 保留：真实动物照片（活体、自然姿态、野外或自然环境）
✗ 删除：标本照片、插画、图表、文字图片、logo、标志、空场景、纯风景

只回答：KEEP 或 DELETE"""


def encode_image_to_base64(path: str) -> str:
    """读取图片并压缩后转为 base64"""
    img = Image.open(path).convert("RGB")  # 转为 RGB 避免 RGBA 问题
    # 缩放图片以减少 API 负载
    img.thumbnail((MAX_IMAGE_SIZE, MAX_IMAGE_SIZE), Image.Resampling.LANCZOS)
    buffer = BytesIO()
    img.save(buffer, format="JPEG", quality=70)
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def call_localai(image_path: str, prompt: str) -> str:
    """调用 LocalAI 检测图片是否包含活体动物"""
    image_b64 = encode_image_to_base64(image_path)

    data = {
        "model": MODEL_NAME,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
                ],
            }
        ],
        "temperature": 0.1,
    }

    for attempt in range(3):
        try:
            with urllib.request.urlopen(LOCAL_AI_URL, data=json.dumps(data).encode("utf-8"), timeout=60) as resp:
                if resp.status == 200:
                    result = json.loads(resp.read().decode("utf-8"))
                    return result.get("choices", [{}])[0].get("message", {}).get("content", "").strip().upper()
        except Exception as e:
            if attempt < 2:
                import time
                time.sleep(2 * (attempt + 1))
            else:
                print(f"  ! API 错误: {e}", flush=True)
    return "ERROR"


def main():
    total_deleted = 0

    for cls in CLASSES:
        cls_dir = os.path.join(IMG_DIR, cls)
        if not os.path.exists(cls_dir):
            continue

        files = sorted(f for f in os.listdir(cls_dir) if f.endswith(".jpg"))
        print(f"\n[filter] {cls}: {len(files)} 张", flush=True)

        deleted = 0
        kept = 0

        for f in files:
            path = os.path.join(cls_dir, f)
            prompt = FILTER_PROMPT.format(target=cls)
            result = call_localai(path, prompt)

            if "DELETE" in result:
                os.remove(path)
                deleted += 1
                print(f"  {f}: DELETE", flush=True)
            else:
                kept += 1
                print(f"  {f}: KEEP", flush=True)

        print(f"  → 保留 {kept} 张，删除 {deleted} 张", flush=True)
        total_deleted += deleted

    print(f"\n[done] 共删除 {total_deleted} 张图片", flush=True)


if __name__ == "__main__":
    sys.exit(main())