Update

app.py

@@ -0,0 +1,166 @@
+import gradio as gr
+import csv
+import os
+from pathlib import Path
+
+import cv2
+import numpy as np
+from PIL import Image
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+import spaces
+
+# 画像のサイズ設定
+IMAGE_SIZE = 448
+
+def preprocess_image(image):
+ image = np.array(image)
+ image = image[:, :, ::-1] # BGRからRGBへ変換
+
+ # 画像を正方形にするためのパディングを追加
+ size = max(image.shape[0:2])
+ pad_x = size - image.shape[1]
+ pad_y = size - image.shape[0]
+ pad_l = pad_x // 2
+ pad_t = pad_y // 2
+ image = np.pad(image, ((pad_t, pad_y - pad_t), (pad_l, pad_x - pad_l), (0, 0)), mode="constant", constant_values=255)
+
+ # サイズに合わせた補間方法を選択
+ interp = cv2.INTER_AREA if size > IMAGE_SIZE else cv2.INTER_LANCZOS4
+ image = cv2.resize(image, (IMAGE_SIZE, IMAGE_SIZE), interpolation=interp)
+ image = image.astype(np.float32)
+ return image
+
+@spaces.GPU
+def process_image(image_path, input_name, ort_sess, rating_tags, character_tags, general_tags):
+ thresh = 0.35
+ try:
+ image = Image.open(image_path)
+ image = image.convert("RGB") if image.mode != "RGB" else image
+ image = preprocess_image(image)
+ except Exception as e:
+ print(f"画像を読み込めません: {image_path}, エラー: {e}")
+ return
+
+ img = np.array([image])
+ prob = ort_sess.run(None, {input_name: img})[0][0] # ONNXモデルからの出力
+
+ # NSFW/SFW判定
+ tag_confidences = {tag: prob[i] for i, tag in enumerate(rating_tags)}
+ max_nsfw_score = max(tag_confidences.get("questionable", 0), tag_confidences.get("explicit", 0))
+ max_sfw_score = tag_confidences.get("general", 0)
+ NSFW_flag = None
+
+ if max_nsfw_score > max_sfw_score:
+ NSFW_flag = "NSFWの可能性が高いです"
+ else:
+ NSFW_flag = "SFWの可能性が高いです"
+
+ # 版権キャラクターの可能性を評価
+ character_tags_with_probs = []
+ for i, p in enumerate(prob[4:]):
+ if p >= thresh and i >= len(general_tags):
+ tag_index = i - len(general_tags)
+ if tag_index < len(character_tags):
+ tag_name = character_tags[tag_index]
+ prob_percent = round(p * 100, 2) # 確率をパーセンテージに変換
+ character_tags_with_probs.append((tag_name, f"{prob_percent}%"))
+
+ IP_flag = None
+ if character_tags_with_probs:
+ IP_flag = f"版権キャラクター: {character_tags_with_probs}の可能性があります"
+ else:
+ IP_flag = "版権キャラクターの可能性が低いと思われます"
+
+ # タグを生成
+ tag_freq = {}
+ undesired_tags = [] 
+ combined_tags = []
+ general_tag_text = ""
+ character_tag_text = ""
+ remove_underscore = True
+ caption_separator = ", "
+ general_threshold = 0.35
+ character_threshold = 0.35
+
+ for i, p in enumerate(prob[4:]):
+ if i < len(general_tags) and p >= general_threshold:
+ tag_name = general_tags[i]
+ if remove_underscore and len(tag_name) > 3: # ignore emoji tags like >_< and ^_^
+ tag_name = tag_name.replace("_", " ")
+
+ if tag_name not in undesired_tags:
+ tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
+ general_tag_text += caption_separator + tag_name
+ combined_tags.append(tag_name)
+ elif i >= len(general_tags) and p >= character_threshold:
+ tag_name = character_tags[i - len(general_tags)]
+ if remove_underscore and len(tag_name) > 3:
+ tag_name = tag_name.replace("_", " ")
+
+ if tag_name not in undesired_tags:
+ tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
+ character_tag_text += caption_separator + tag_name
+ combined_tags.append(tag_name)
+
+ # 先頭のカンマを取る
+ if len(general_tag_text) > 0:
+ general_tag_text = general_tag_text[len(caption_separator) :]
+ if len(character_tag_text) > 0:
+ character_tag_text = character_tag_text[len(caption_separator) :]
+ tag_text = caption_separator.join(combined_tags)
+ 
+ return NSFW_flag, IP_flag, tag_text
+
+
+class webui:
+ def __init__(self):
+ self.demo = gr.Blocks()
+
+ @spaces.GPU
+ def main(self, image_path, model_id):
+ print("Hugging Faceからモデルをダウンロード中")
+ onnx_path = hf_hub_download(model_id, "model.onnx")
+ csv_path = hf_hub_download(model_id, "selected_tags.csv")
+
+ print("ONNXモデルを実行中")
+ print(f"ONNXモデルのパス: {onnx_path}")
+
+ ort_sess = ort.InferenceSession(onnx_path)
+
+ with open(csv_path, "r", encoding="utf-8") as f:
+ reader = csv.reader(f)
+ header = next(reader)
+ rows = list(reader)
+ assert header == ["tag_id", "name", "category", "count"], f"CSVフォーマット���期待と異なります: {header}"
+
+ rating_tags = [row[1] for row in rows if row[2] == "9"]
+ character_tags = [row[1] for row in rows if row[2] == "4"]
+ general_tags = [row[1] for row in rows[1:] if row[2] == "0"]
+
+ NSFW_flag, IP_flag, tag_text = process_image(image_path, ort_sess.get_inputs()[0].name, ort_sess, rating_tags, character_tags, general_tags)
+ return NSFW_flag, IP_flag, tag_text
+
+
+ def launch(self):
+ with self.demo:
+ with gr.Row():
+ with gr.Column():
+ input_image = gr.Image(type='filepath', label="Analysis Image")
+ model_id = gr.Textbox(label="NSFW Flag", value="SmilingWolf/wd-vit-tagger-v3")
+ output_0 = gr.Textbox(label="NSFW Flag")
+ output_1 = gr.Textbox(label="IP Flag")
+ output_2 = gr.Textbox(label="Tags")
+ submit = gr.Button(value="Start Analysis")
+ 
+ submit.click(
+ self.main, 
+ inputs=[input_image, model_id], 
+ outputs=[output_0, output_1, output_2]
+ )
+
+ self.demo.launch()
+
+if __name__ == "__main__":
+ ui = webui()
+ ui.launch()

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python
+numpy
+Pillow
+onnxruntime
+onnxruntime-gpu