Update 0filterImage.py

0filterImage.py

@@ -7,32 +7,22 @@ import os
 from io import BytesIO
 import pandas as pd
 from PIL import Image
-
+from dotenv import load_dotenv
 import requests
+from transformers import pipeline
 
 
-def ocr(image):
-
-    image = Image.open(image)
-    img_buffer = BytesIO()
-    image.save(img_buffer, format=image.format)
-    byte_data = img_buffer.getvalue()
-    base64_bytes = base64.b64encode(byte_data)  # bytes
-    base64_str = base64_bytes.decode()
-    url = "https://www.modelscope.cn/api/v1/studio/damo/ofa_ocr_pipeline/gradio/api/predict/"
-    payload = json.dumps({
-        "data": [f"data:image/jpeg;base64,{base64_str}"],
-        "dataType": ["image"]
-    })
-    headers = {
-        'Content-Type': 'application/json'
-    }
-
-    response = requests.request("POST", url, headers=headers, data=payload)
-    jobj = json.loads(response.text)
-    return jobj
+def get_nsfw_score(image_path:str,model:"模型")->float:
+    #输入图片和模型,返回是否有问题
+    img = Image.open(image_path)
+    result = model(images=img)
+    nsfw_score = next((item['score'] for item in result if item['label']=='nsfw'),None)
+    return nsfw_score
+    
 
 if __name__ == '__main__':
+    load_dotenv()
+    model = pipeline("image-classification", model="Falconsai/nsfw_image_detection")#加载模型
     # 获取当前目录的子目录的路径
     img_path = 'manga'
     subdir_path = os.path.join(os.getcwd(), img_path)
@@ -44,8 +34,8 @@ if __name__ == '__main__':
             if file.endswith(".jpg") or file.endswith(".png"):
                 image_files.append(os.path.relpath(os.path.join(root, file)))
     for image_path in image_files:
-        result = ocr(image_path) ##dataframe格式  有两列  boxid 和text
-        if 'error' in result:
+        result = get_nsfw_score(image_path)#返回float的得分
+        if result> 0.5:
             print("发现问题图片,需要删除以过审:",image_path)
             os.remove(image_path)
         else: