Update 0filterImage.py

0filterImage.py

@@ -1,28 +1,182 @@
-#通过modelscope接口对问题图片予以删除，保证过审
-
-
-import base64
-import json
-import os
-from io import BytesIO
+import numpy as np
+import os, re, cv2
+from typing import *
 import pandas as pd
 from PIL import Image
-from dotenv import load_dotenv
-import requests
-from transformers import pipeline
+from huggingface_hub import hf_hub_download
+from onnxruntime import InferenceSession
+
+
+
+# noinspection PyUnresolvedReferences
+def make_square(img, target_size):
+    old_size = img.shape[:2]
+    desired_size = max(old_size)
+    desired_size = max(desired_size, target_size)
+
+    delta_w = desired_size - old_size[1]
+    delta_h = desired_size - old_size[0]
+    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
+    left, right = delta_w // 2, delta_w - (delta_w // 2)
+
+    color = [255, 255, 255]
+    return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
+
+
+# noinspection PyUnresolvedReferences
+def smart_resize(img, size):
+    # Assumes the image has already gone through make_square
+    if img.shape[0] > size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+    elif img.shape[0] < size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
+    else:  # just do nothing
+        pass
+
+    return img
+
+
+class WaifuDiffusionInterrogator:
+    def __init__(
+            self,
+            repo='SmilingWolf/wd-v1-4-vit-tagger',
+            model_path='model.onnx',
+            tags_path='selected_tags.csv',
+            mode: str = "auto"
+    ) -> None:
+        self.__repo = repo
+        self.__model_path = model_path
+        self.__tags_path = tags_path
+        self._provider_mode = mode
+
+        self.__initialized = False
+        self._model, self._tags = None, None
+
+    def _init(self) -> None:
+        if self.__initialized:
+            return
+
+        model_path = hf_hub_download(self.__repo, filename=self.__model_path)
+        tags_path = hf_hub_download(self.__repo, filename=self.__tags_path)
+
+        self._model = InferenceSession(str(model_path))
+        self._tags = pd.read_csv(tags_path)
+
+        self.__initialized = True
+
+    def _calculation(self, image: Image.Image) -> pd.DataFrame:
+        # print(image) todo: figure out what to do if URL
+        self._init()
+
+        # code for converting the image and running the model is taken from the link below
+        # thanks, SmilingWolf!
+        # https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags/blob/main/app.py
+
+        # convert an image to fit the model
+        _, height, _, _ = self._model.get_inputs()[0].shape
+
+        # alpha to white
+        print(image)
+        image = image.convert('RGBA')
+        new_image = Image.new('RGBA', image.size, 'WHITE')
+        new_image.paste(image, mask=image)
+        image = new_image.convert('RGB')
+        image = np.asarray(image)
+
+        # PIL RGB to OpenCV BGR
+        image = image[:, :, ::-1]
+
+        image = make_square(image, height)
+        image = smart_resize(image, height)
+        image = image.astype(np.float32)
+        image = np.expand_dims(image, 0)
+
+        # evaluate model
+        input_name = self._model.get_inputs()[0].name
+        label_name = self._model.get_outputs()[0].name
+        confidence = self._model.run([label_name], {input_name: image})[0]
+
+        full_tags = self._tags[['name', 'category']].copy()
+        full_tags['confidence'] = confidence[0]
+
+        return full_tags
+
+    def interrogate(self, image: Image) -> Tuple[Dict[str, float], Dict[str, float]]:
+        full_tags = self._calculation(image)
+
+        # first 4 items are for rating (general, sensitive, questionable, explicit)
+        ratings = dict(full_tags[full_tags['category'] == 9][['name', 'confidence']].values)
+
+        # rest are regular tags
+        tags = dict(full_tags[full_tags['category'] != 9][['name', 'confidence']].values)
+
+        return ratings, tags
+
+
+WAIFU_MODELS: Mapping[str, WaifuDiffusionInterrogator] = {
+    'chen-vit': WaifuDiffusionInterrogator(),
+    'chen-convnext': WaifuDiffusionInterrogator(
+        repo='SmilingWolf/wd-v1-4-convnext-tagger'
+    ),
+    'chen-convnext2': WaifuDiffusionInterrogator(
+        repo="SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
+    ),
+    'chen-swinv2': WaifuDiffusionInterrogator(
+        repo='SmilingWolf/wd-v1-4-swinv2-tagger-v2'
+    ),
+    'chen-moat2': WaifuDiffusionInterrogator(
+        repo='SmilingWolf/wd-v1-4-moat-tagger-v2'
+    ),
+    'chen-convnext3': WaifuDiffusionInterrogator(
+        repo='SmilingWolf/wd-convnext-tagger-v3'
+    ),
+    'chen-vit3': WaifuDiffusionInterrogator(
+        repo='SmilingWolf/wd-vit-tagger-v3'
+    ),
+    'chen-swinv3': WaifuDiffusionInterrogator(
+        repo='SmilingWolf/wd-swinv2-tagger-v3'
+    ),
+}
+RE_SPECIAL = re.compile(r'([\\()])')
+
+
+def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float,
+                       use_spaces: bool, use_escape: bool, include_ranks=False, score_descend=True) \
+        -> Tuple[Mapping[str, float], str, Mapping[str, float]]:
+    model = WAIFU_MODELS[model_name]
+    ratings, tags = model.interrogate(image)
+
+    filtered_tags = {
+        tag: score for tag, score in tags.items()
+        if score >= threshold
+    }
+
+    text_items = []
+    tags_pairs = filtered_tags.items()
+    if score_descend:
+        tags_pairs = sorted(tags_pairs, key=lambda x: (-x[1], x[0]))
+    for tag, score in tags_pairs:
+        tag_outformat = tag
+        if use_spaces:
+            tag_outformat = tag_outformat.replace('_', '-')
+        else:
+            tag_outformat = tag_outformat.replace(' ', ', ')
+            tag_outformat = tag_outformat.replace('_', ' ')
+        if use_escape:
+            tag_outformat = re.sub(RE_SPECIAL, r'\\\1', tag_outformat)
+        if include_ranks:
+            tag_outformat = f"({tag_outformat}:{score:.3f})"
+        text_items.append(tag_outformat)
+    if use_spaces:
+        output_text = ' '.join(text_items)
+    else:
+        output_text = ', '.join(text_items)
+
+    return ratings, output_text, filtered_tags
 
 
-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)
@@ -34,11 +188,20 @@ 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 = get_nsfw_score(image_path,model)#返回float的得分
-        if result> 0.5:
-            print("发现问题图片,需要删除以过审:",image_path)
+        # 打开并读取图像文件
+        image_data = Image.open(image_path)
+        result = image_to_wd14_tags(image_data, 'chen-moat2', 0.5, True, True)#传入数据判断标签,然后只看rating tag就行，即第[0]个
+        # 从 result 中提取第一个元素(rating)
+        rating_dict = result[0]
+        # 找到占比最大的元素
+        max_proportion_key = max(rating_dict, key=rating_dict.get)
+        max_proportion_value = rating_dict[max_proportion_key]
+
+        # 输出占比最大的元素
+        print(f"占比最大的元素为：{max_proportion_key}，占比为：{max_proportion_value}")
+        if max_proportion_key=="questionable" or max_proportion_key=="explicit":
+            print("图片不合格，开始删除")
             os.remove(image_path)
+            print("成功删除不合格图片")
         else:
-            print(image_path, "图片没有问题")
-
-
+            print("图片合格")