init

.gitignore

@@ -1,2 +1,5 @@
 *.jpg
-*.png
+*.png
+hf_cache/
+models/
+__pycache__/

app.py

@@ -1,10 +1,19 @@
 import gradio as gr
 import threading
-from util import process_image_edit, check_nsfw, get_country_info_safe
+from util import process_image_edit, get_country_info_safe
+from nfsw import NSFWDetector
 
 IP_Dict = {}
 NSFW_Dict = {}  # 记录每个IP的NSFW违规次数
 
+# 初始化NSFW检测器（从Hugging Face下载）
+try:
+    nsfw_detector = NSFWDetector()  # 自动从Hugging Face下载falconsai_yolov9_nsfw_model_quantized.pt
+    print("✅ NSFW检测器初始化成功")
+except Exception as e:
+    print(f"❌ NSFW检测器初始化失败: {e}")
+    nsfw_detector = None
+
 def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.Progress()):
     """
     Interface function for processing image editing
@@ -24,7 +33,7 @@ def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.P
     # 检查IP是否因NSFW违规过多而被屏蔽 3 
     if client_ip in NSFW_Dict and NSFW_Dict[client_ip] >= 3:
         print(f"❌ IP blocked due to excessive NSFW violations - IP: {client_ip}({country_info}), violations: {NSFW_Dict[client_ip]}")
-        return None, "❌ NSFW content too much. Access denied due to policy violations"
+        return None, f"❌ Your ip {client_ip}，your region has been blocked"
     
     if input_image is None:
         return None, "Please upload an image first"
@@ -36,14 +45,21 @@ def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.P
     if len(prompt.strip()) <= 3:
         return None, "❌ Editing prompt must be more than 3 characters"
     
-    # 检查是否包含NSFW内容
-    if check_nsfw(prompt.strip()) == 1:
-        # 记录NSFW违规次数
-        if client_ip not in NSFW_Dict:
-            NSFW_Dict[client_ip] = 0
-        NSFW_Dict[client_ip] += 1
-        print(f"❌ NSFW content detected - IP: {client_ip}({country_info}), violations: {NSFW_Dict[client_ip]}, prompt: {prompt.strip()}")
-        return None, "❌ NSFW content detected. Please modify your prompt. NSFW detected {NSFW_Dict[client_ip]}/10 times"
+    # 检查图片是否包含NSFW内容
+    nsfw_result = None
+    if nsfw_detector is not None:
+        try:
+            nsfw_result = nsfw_detector.predict_label_only(input_image)
+            if nsfw_result.lower() == "nsfw":
+                # 记录NSFW违规次数
+                if client_ip not in NSFW_Dict:
+                    NSFW_Dict[client_ip] = 0
+                NSFW_Dict[client_ip] += 1
+                print(f"❌ NSFW image detected - IP: {client_ip}({country_info}), violations: {NSFW_Dict[client_ip]}")
+                return None, f"❌ Your ip {client_ip}，your region has been blocked"
+        except Exception as e:
+            print(f"⚠️ NSFW检测失败: {e}")
+            # 检测失败时允许继续处理
     
     if IP_Dict[client_ip]>8 and country_info.lower() in ["印度", "巴基斯坦"]:
         print(f"❌ Content not allowed - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
@@ -67,17 +83,17 @@ def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.P
     
     try:
         # 打印成功访问的信息
-        print(f"✅ Processing started - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
+        print(f"✅ Processing started - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}", flush=True)
         
         # Call image editing processing function
         result_url, message = process_image_edit(input_image, prompt.strip(), progress_callback)
         
         if result_url:
-            print(f"✅ Processing completed successfully - IP: {client_ip}({country_info}), result_url: {result_url}")
+            print(f"✅ Processing completed successfully - IP: {client_ip}({country_info}), result_url: {result_url}", flush=True)
             progress(1.0, desc="Processing completed")
             return result_url, "✅ " + message
         else:
-            print(f"❌ Processing failed - IP: {client_ip}({country_info}), error: {message}")
+            print(f"❌ Processing failed - IP: {client_ip}({country_info}), error: {message}", flush=True)
             return None, "❌ " + message
             
     except Exception as e:

labels.json

@@ -0,0 +1,4 @@
+{
+    "0": "normal",
+    "1": "nsfw"
+}

nfsw.py

@@ -0,0 +1,210 @@
+import os
+from PIL import Image
+import numpy as np
+import onnxruntime as ort
+import json
+from huggingface_hub import hf_hub_download
+
+
+class NSFWDetector:
+    """
+    NSFW检测器类，使用YOLOv9模型进行图像分类
+    """
+    
+    def __init__(self, repo_id="Falconsai/nsfw_image_detection", 
+                 model_filename="falconsai_yolov9_nsfw_model_quantized.pt",
+                 labels_filename="labels.json",
+                 input_size=(224, 224)):
+        """
+        初始化NSFW检测器
+        
+        Args:
+            repo_id (str): Hugging Face仓库ID
+            model_filename (str): 模型文件名
+            labels_filename (str): 标签文件名
+            input_size (tuple): 模型输入尺寸 (height, width)
+        """
+        self.repo_id = repo_id
+        self.model_filename = model_filename
+        self.labels_filename = labels_filename
+        self.input_size = input_size
+        
+        # 从Hugging Face下载文件
+        self.model_path = self._download_model()
+        self.labels_path = self._download_labels()
+        
+        # 加载标签
+        self.labels = self._load_labels()
+        
+        # 加载模型
+        self.session = self._load_model()
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_name = self.session.get_outputs()[0].name
+        
+    def _download_model(self):
+        """
+        从Hugging Face下载模型文件
+        
+        Returns:
+            str: 下载的模型文件路径
+        """
+        try:
+            print(f"正在从 {self.repo_id} 下载模型文件: {self.model_filename}")
+            model_path = hf_hub_download(
+                repo_id=self.repo_id,
+                filename=self.model_filename,
+                cache_dir="./hf_cache"
+            )
+            print(f"✅ 模型下载成功: {model_path}")
+            return model_path
+        except Exception as e:
+            raise RuntimeError(f"模型下载失败: {e}")
+    
+    def _download_labels(self):
+        """
+        从Hugging Face下载标签文件
+        
+        Returns:
+            str: 下载的标签文件路径
+        """
+        try:
+            print(f"正在从 {self.repo_id} 下载标签文件: {self.labels_filename}")
+            labels_path = hf_hub_download(
+                repo_id=self.repo_id,
+                filename=self.labels_filename,
+                cache_dir="./hf_cache"
+            )
+            print(f"✅ 标签文件下载成功: {labels_path}")
+            return labels_path
+        except Exception as e:
+            raise RuntimeError(f"标签文件下载失败: {e}")
+    
+    def _load_labels(self):
+        """
+        加载类别标签
+        
+        Returns:
+            dict: 标签字典
+        """
+        try:
+            with open(self.labels_path, "r") as f:
+                return json.load(f)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"标签文件未找到: {self.labels_path}")
+        except json.JSONDecodeError:
+            raise ValueError(f"标签文件格式错误: {self.labels_path}")
+    
+    def _load_model(self):
+        """
+        加载ONNX模型
+        
+        Returns:
+            onnxruntime.InferenceSession: 模型会话
+        """
+        try:
+            return ort.InferenceSession(self.model_path)
+        except Exception as e:
+            raise RuntimeError(f"模型加载失败: {self.model_path}, 错误: {e}")
+    
+    def _preprocess_image(self, image_path):
+        """
+        图像预处理
+        
+        Args:
+            image_path (str): 图像文件路径
+            
+        Returns:
+            tuple: (预处理后的张量, 原始图像)
+        """
+        try:
+            # 加载并转换图像
+            original_image = Image.open(image_path).convert("RGB")
+            
+            # 调整尺寸
+            image_resized = original_image.resize(self.input_size, Image.Resampling.BILINEAR)
+            
+            # 转换为numpy数组并归一化
+            image_np = np.array(image_resized, dtype=np.float32) / 255.0
+            
+            # 调整维度顺序 [H, W, C] -> [C, H, W]
+            image_np = np.transpose(image_np, (2, 0, 1))
+            
+            # 添加批次维度 [C, H, W] -> [1, C, H, W]
+            input_tensor = np.expand_dims(image_np, axis=0).astype(np.float32)
+            
+            return input_tensor, original_image
+            
+        except FileNotFoundError:
+            raise FileNotFoundError(f"图像文件未找到: {image_path}")
+        except Exception as e:
+            raise RuntimeError(f"图像预处理失败: {e}")
+    
+    def _postprocess_predictions(self, predictions):
+        """
+        后处理预测结果
+        
+        Args:
+            predictions: 模型预测输出
+            
+        Returns:
+            str: 预测的类别标签
+        """
+        predicted_index = np.argmax(predictions)
+        predicted_label = self.labels[str(predicted_index)]
+        return predicted_label
+    
+    def predict(self, image_path):
+        """
+        对单张图像进行NSFW检测
+        
+        Args:
+            image_path (str): 图像文件路径
+            
+        Returns:
+            tuple: (预测标签, 原始图像)
+        """
+        # 预处理图像
+        input_tensor, original_image = self._preprocess_image(image_path)
+        
+        # 运行推理
+        outputs = self.session.run([self.output_name], {self.input_name: input_tensor})
+        predictions = outputs[0]
+        
+        # 后处理结果
+        predicted_label = self._postprocess_predictions(predictions)
+        
+        return predicted_label, original_image
+    
+    def predict_label_only(self, image_path):
+        """
+        只返回预测标签（不返回图像）
+        
+        Args:
+            image_path (str): 图像文件路径
+            
+        Returns:
+            str: 预测的类别标签
+        """
+        predicted_label, _ = self.predict(image_path)
+        return predicted_label
+
+# --- 使用示例 ---
+if __name__ == "__main__":
+    # 配置参数
+    single_image_path = "datas/bad01.jpg"
+    
+    try:
+        # 创建检测器实例（自动从Hugging Face下载）
+        detector = NSFWDetector()
+        
+        # 检查图像文件是否存在
+        if os.path.exists(single_image_path):
+            # 进行预测
+            predicted_label = detector.predict_label_only(single_image_path)
+            print(f"图像文件: {single_image_path}")
+            print(f"预测结果: {predicted_label}")
+        else:
+            print(f"错误: 指定的图像文件不存在: {single_image_path}")
+            
+    except Exception as e:
+        print(f"初始化检测器时发生错误: {e}")

requirements.txt

@@ -4,4 +4,7 @@ requests>=2.28.0
 func-timeout>=4.3.5
 numpy>=1.24.0
 boto3
-botocore
+botocore
+onnxruntime
+huggingface_hub>=0.16.0
+Pillow>=9.0.0

util.py

@@ -177,44 +177,6 @@ def get_country_info_safe(ip):
         return "Unknown"
 
 
-def check_nsfw(prompt):
-    """
-    检查prompt是否包含NSFW内容，包含返回1，否则返回0
-    """
-    try:
-        response = requests.post(
-            url="https://openrouter.ai/api/v1/chat/completions",
-            headers={
-                "Authorization": f"Bearer {LLMKEY}",
-                "Content-Type": "application/json",
-            },
-            data=json.dumps({
-                "model": "google/gemini-2.5-flash",
-                "messages": [
-                    {
-                        "role": "system",
-                        "content": "你是一个nsfw指令判断助手，请判断用户输入的prompt指令是否会导致nsfw内容? 你只需要回答 是 或者 否"
-                    },
-                    {
-                        "role": "user",
-                        "content": prompt
-                    }
-                ],
-            })
-        )
-        res_json = response.json()
-        # 兼容不同模型返回格式
-        if "choices" in res_json and len(res_json["choices"]) > 0:
-            content = res_json["choices"][0].get("message", {}).get("content", "")
-            if "是" in content:
-                return 1
-            else:
-                return 0
-        else:
-            return 0
-    except Exception as e:
-        # 出错时默认返回0
-        return 0
 
 
 def submit_image_edit_task(user_image_url, prompt):