Upload 2 files

filter_nsfw.py

@@ -0,0 +1,83 @@
+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()