init

app.py

@@ -1,5 +1,8 @@
 import gradio as gr
 import threading
+import os
+import shutil
+import tempfile
 from util import process_image_edit, get_country_info_safe
 from nfsw import NSFWDetector
 
@@ -31,9 +34,9 @@ def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.P
     country_info = get_country_info_safe(client_ip)
     
     # 检查IP是否因NSFW违规过多而被屏蔽 3 
-    if client_ip in NSFW_Dict and NSFW_Dict[client_ip] >= 3:
+    if client_ip in NSFW_Dict and NSFW_Dict[client_ip] >= 5:
         print(f"❌ IP blocked due to excessive NSFW violations - IP: {client_ip}({country_info}), violations: {NSFW_Dict[client_ip]}")
-        return None, f"❌ Your ip {client_ip}，your region has been blocked"
+        return None, f"❌ Your ip {client_ip}，your region has been blocked for too much nsfw content"
     
     if input_image is None:
         return None, "Please upload an image first"
@@ -47,31 +50,34 @@ def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.P
     
     # 检查图片是否包含NSFW内容
     nsfw_result = None
-    if nsfw_detector is not None:
+    if nsfw_detector is not None and input_image is not None:
         try:
-            nsfw_result = nsfw_detector.predict_label_only(input_image)
+            # 直接使用PIL Image对象进行检测，避免文件路径问题
+            nsfw_result = nsfw_detector.predict_pil_label_only(input_image)
+            print(f"🔍 NSFW检测结果: {nsfw_result} - IP: {client_ip}({country_info})")
+            
             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"
+                return None, f"❌ Your ip {client_ip}，your region has been blocked for too much nsfw content"
+                
         except Exception as e:
             print(f"⚠️ NSFW检测失败: {e}")
             # 检测失败时允许继续处理
     
-    if IP_Dict[client_ip]>8 and country_info.lower() in ["印度", "巴基斯坦"]:
+    if IP_Dict[client_ip]>10 and country_info.lower() in ["印度", "巴基斯坦"]:
         print(f"❌ Content not allowed - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
         return None, "❌ Content not allowed. Please modify your prompt"
-    if IP_Dict[client_ip]>18 and country_info.lower() in ["中国"]:
-        print(f"❌ Content not allowed - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
-        return None, "❌ Content not allowed. Please modify your prompt"
-        
-    if client_ip.lower() in ["221.194.171.230", "101.126.56.37", "101.126.56.44"]:
-        print(f"❌ Content not allowed - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
-        return None, "❌ Content not allowed. Please modify your prompt"
-        
+    # if IP_Dict[client_ip]>18 and country_info.lower() in ["中国"]:
+    #     print(f"❌ Content not allowed - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
+    #     return None, "❌ Content not allowed. Please modify your prompt"
+    # if client_ip.lower() in ["221.194.171.230", "101.126.56.37", "101.126.56.44"]:
+    #     print(f"❌ Content not allowed - IP: {client_ip}({country_info}), count: {IP_Dict[client_ip]}, prompt: {prompt.strip()}")
+    #     return None, "❌ Content not allowed. Please modify your prompt"
+    
 
     result_url = None
     status_message = ""
@@ -136,7 +142,7 @@ def create_app():
                 gr.Markdown("### 📸 Upload Image")
                 input_image = gr.Image(
                     label="Select image to edit",
-                    type="filepath",
+                    type="pil",
                     height=400,
                     elem_classes=["upload-area"]
                 )

nfsw.py

@@ -187,6 +187,58 @@ class NSFWDetector:
         """
         predicted_label, _ = self.predict(image_path)
         return predicted_label
+    
+    def predict_from_pil(self, pil_image):
+        """
+        直接从PIL Image对象进行NSFW检测
+        
+        Args:
+            pil_image (PIL.Image): PIL图像对象
+            
+        Returns:
+            tuple: (预测标签, 原始图像)
+        """
+        try:
+            # 确保是RGB格式
+            if pil_image.mode != "RGB":
+                pil_image = pil_image.convert("RGB")
+            
+            # 调整尺寸
+            image_resized = pil_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)
+            
+            # 运行推理
+            outputs = self.session.run([self.output_name], {self.input_name: input_tensor})
+            predictions = outputs[0]
+            
+            # 后处理结果
+            predicted_label = self._postprocess_predictions(predictions)
+            
+            return predicted_label, pil_image
+            
+        except Exception as e:
+            raise RuntimeError(f"PIL图像预测失败: {e}")
+    
+    def predict_pil_label_only(self, pil_image):
+        """
+        从PIL Image对象只返回预测标签
+        
+        Args:
+            pil_image (PIL.Image): PIL图像对象
+            
+        Returns:
+            str: 预测的类别标签
+        """
+        predicted_label, _ = self.predict_from_pil(pil_image)
+        return predicted_label
 
 # --- 使用示例 ---
 if __name__ == "__main__":

util.py

@@ -11,7 +11,9 @@ import func_timeout
 import numpy as np
 import gradio as gr
 import boto3
+import tempfile
 from botocore.client import Config
+from PIL import Image
 
 
 # TOKEN = os.environ['TOKEN']
@@ -248,15 +250,38 @@ def check_task_status(task_id):
         return 'error', None, f"请求异常: {str(e)}"
 
 
-def process_image_edit(img_path, prompt, progress_callback=None):
+def process_image_edit(img_input, prompt, progress_callback=None):
     """
     处理图片编辑的完整流程
+    
+    Args:
+        img_input: 可以是文件路径(str)或PIL Image对象
+        prompt: 编辑指令
+        progress_callback: 进度回调函数
     """
+    temp_img_path = None
     try:
         # 生成客户端 IP 和时间戳
         client_ip = "127.0.0.1"  # 默认IP
         time_id = int(time.time())
         
+        # 处理输入图像 - 支持PIL Image和文件路径
+        if hasattr(img_input, 'save'):  # PIL Image对象
+            # 创建临时文件
+            temp_dir = tempfile.mkdtemp()
+            temp_img_path = os.path.join(temp_dir, f"temp_img_{time_id}.jpg")
+            
+            # 保存PIL Image为临时文件
+            if img_input.mode != 'RGB':
+                img_input = img_input.convert('RGB')
+            img_input.save(temp_img_path, 'JPEG', quality=95)
+            
+            img_path = temp_img_path
+            print(f"💾 PIL Image已保存为临时文件: {temp_img_path}")
+        else:
+            # 假设是文件路径
+            img_path = img_input
+        
         if progress_callback:
             progress_callback("uploading image...")
         
@@ -305,6 +330,19 @@ def process_image_edit(img_path, prompt, progress_callback=None):
         
     except Exception as e:
         return None, f"error occurred during processing: {str(e)}"
+    
+    finally:
+        # 清理临时文件
+        if temp_img_path and os.path.exists(temp_img_path):
+            try:
+                os.remove(temp_img_path)
+                # 尝试删除临时目录（如果为空）
+                temp_dir = os.path.dirname(temp_img_path)
+                if os.path.exists(temp_dir):
+                    os.rmdir(temp_dir)
+                print(f"🗑️ 已清理临时文件: {temp_img_path}")
+            except Exception as cleanup_error:
+                print(f"⚠️ 清理临时文件失败: {cleanup_error}")
 
 
 if __name__ == "__main__":