loras_tools / filter_nsfw.py

Upload 2 files

705e0f6 verified about 21 hours ago

2.93 kB

import os
import shutil
import glob
from nudenet import NudeDetector
from PIL import Image, UnidentifiedImageError

def main():
    current_dir = os.getcwd()
    nsfw_dir = os.path.join(current_dir, "NSFW")
    
    # 1. 创建 NSFW 文件夹
    if not os.path.exists(nsfw_dir):
        os.makedirs(nsfw_dir)
        print(f"📁 已创建隔离文件夹: {nsfw_dir}")

    # 2. 初始化检测器 (首次运行会自动下载模型)
    print("⏳ 正在加载 NSFW 检测模型，首次运行可能需要下载...")
    detector = NudeDetector()
    
    # 定义属于 NSFW 的标签 (根据 nudenet 的标签库)
    NSFW_LABELS = {
        "FEMALE_BREAST_EXPOSED",
        "FEMALE_GENITALIA_EXPOSED",
        "MALE_GENITALIA_EXPOSED",
        "ANUS_EXPOSED",
        "BUTTOCKS_EXPOSED",
        # 加入以下擦边标签：
        "MALE_BREAST_EXPOSED",      # 男士露胸（有时候会被误判）
    }

    # 3. 查找当前目录下所有的 jpg
    jpg_files = glob.glob(os.path.join(current_dir, "*.jpg")) + glob.glob(os.path.join(current_dir, "*.jpeg"))
    
    if not jpg_files:
        print("❌ 当前目录下没有找到 .jpg 或 .jpeg 文件。")
        return

    print(f"🔍 共找到 {len(jpg_files)} 张图片，开始扫描...\n")

    nsfw_count = 0
    
    for img_path in jpg_files:
        filename = os.path.basename(img_path)
        base_name = os.path.splitext(filename)[0]
        txt_path = os.path.join(current_dir, f"{base_name}.txt")

        try:
            # 检测图片
            detections = detector.detect(img_path)
            
            is_nsfw = False
            found_labels = set()
            
            for det in detections:
                # 如果检测到的标签在 NSFW 列表中，且置信度 > 0.6
                if det["class"] in NSFW_LABELS and det["score"] > 0.6:
                    is_nsfw = True
                    found_labels.add(det["class"])

            if is_nsfw:
                print(f"🚫 发现 NSFW: {filename} -> 特征: {', '.join(found_labels)}")
                
                # 移动图片
                shutil.move(img_path, os.path.join(nsfw_dir, filename))
                
                # 移动同名 TXT (如果有的话)
                if os.path.exists(txt_path):
                    shutil.move(txt_path, os.path.join(nsfw_dir, f"{base_name}.txt"))
                else:
                    print(f"   ⚠️ 警告: {filename} 没有对应的 txt 文件")
                
                nsfw_count += 1
                
        except UnidentifiedImageError:
            print(f"⚠️ 跳过损坏的图片: {filename}")
        except Exception as e:
            print(f"⚠️ 检测 {filename} 时出错: {e}")

    print("\n" + "="*40)
    print(f"🎉 扫描完成！共处理 {len(jpg_files)} 张图，成功隔离 {nsfw_count} 张 NSFW 图片及其文本。")

if __name__ == "__main__":
    main()

	import os
	import shutil
	import glob
	from nudenet import NudeDetector
	from PIL import Image, UnidentifiedImageError

	def main():
	current_dir = os.getcwd()
	nsfw_dir = os.path.join(current_dir, "NSFW")

	# 1. 创建 NSFW 文件夹
	if not os.path.exists(nsfw_dir):
	os.makedirs(nsfw_dir)
	print(f"📁 已创建隔离文件夹: {nsfw_dir}")

	# 2. 初始化检测器 (首次运行会自动下载模型)
	print("⏳ 正在加载 NSFW 检测模型，首次运行可能需要下载...")
	detector = NudeDetector()

	# 定义属于 NSFW 的标签 (根据 nudenet 的标签库)
	NSFW_LABELS = {
	"FEMALE_BREAST_EXPOSED",
	"FEMALE_GENITALIA_EXPOSED",
	"MALE_GENITALIA_EXPOSED",
	"ANUS_EXPOSED",
	"BUTTOCKS_EXPOSED",
	# 加入以下擦边标签：
	"MALE_BREAST_EXPOSED", # 男士露胸（有时候会被误判）
	}

	# 3. 查找当前目录下所有的 jpg
	jpg_files = glob.glob(os.path.join(current_dir, ".jpg")) + glob.glob(os.path.join(current_dir, ".jpeg"))

	if not jpg_files:
	print("❌ 当前目录下没有找到 .jpg 或 .jpeg 文件。")
	return

	print(f"🔍 共找到 {len(jpg_files)} 张图片，开始扫描...\n")

	nsfw_count = 0

	for img_path in jpg_files:
	filename = os.path.basename(img_path)
	base_name = os.path.splitext(filename)[0]
	txt_path = os.path.join(current_dir, f"{base_name}.txt")

	try:
	# 检测图片
	detections = detector.detect(img_path)

	is_nsfw = False
	found_labels = set()

	for det in detections:
	# 如果检测到的标签在 NSFW 列表中，且置信度 > 0.6
	if det["class"] in NSFW_LABELS and det["score"] > 0.6:
	is_nsfw = True
	found_labels.add(det["class"])

	if is_nsfw:
	print(f"🚫 发现 NSFW: {filename} -> 特征: {', '.join(found_labels)}")

	# 移动图片
	shutil.move(img_path, os.path.join(nsfw_dir, filename))

	# 移动同名 TXT (如果有的话)
	if os.path.exists(txt_path):
	shutil.move(txt_path, os.path.join(nsfw_dir, f"{base_name}.txt"))
	else:
	print(f" ⚠️ 警告: {filename} 没有对应的 txt 文件")

	nsfw_count += 1

	except UnidentifiedImageError:
	print(f"⚠️ 跳过损坏的图片: {filename}")
	except Exception as e:
	print(f"⚠️ 检测 {filename} 时出错: {e}")

	print("\n" + "="*40)
	print(f"🎉 扫描完成！共处理 {len(jpg_files)} 张图，成功隔离 {nsfw_count} 张 NSFW 图片及其文本。")

	if __name__ == "__main__":
	main()