113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
"""用 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())
|