Upload 3 files

app.py

@@ -4,10 +4,10 @@ import gradio as gr
 fn.load_model()
 
 with gr.Blocks() as demo:
-    title = gr.Markdown('# Safety Checker')
+    title = gr.Markdown('# Safety Checker 2')
     with gr.Row():
         src_image = gr.Image(label="Source", sources="upload", interactive=True, type="pil")
-        result = gr.Textbox(label="Has NSFW", interactive=False)
+        result = gr.Textbox(label="Result", interactive=False)
 
     src_image.change(
         fn=fn.check,

fn.py

@@ -47,7 +47,7 @@ def check(image):
         image = preprocess_image(image)
     except Exception as e:
         print(f"画像を読み込めません: {image_path}, エラー: {e}")
-        return
+        return None
 
     img = np.array([image])
     prob = ort_sess.run(None, {input_name: img})[0][0]  # ONNXモデルからの出力
@@ -57,11 +57,9 @@ def check(image):
     max_nsfw_score = max(tag_confidences.get("questionable", 0), tag_confidences.get("explicit", 0))
     max_sfw_score = tag_confidences.get("general", 0)
 
-    return max_nsfw_score > max_sfw_score
-
     # 版権キャラクターの可能性を評価
     character_tags_with_probs = []
-    for i, p in enumerate(prob[4:]):
+    for i, p in enumerate(prob[4+len(general_tags):]):
         if p >= thresh and i >= len(general_tags):
             tag_index = i - len(general_tags)
             if tag_index < len(character_tags):
@@ -69,10 +67,21 @@ def check(image):
                 prob_percent = round(p * 100, 2)  # 確率をパーセンテージに変換
                 character_tags_with_probs.append((tag_name, f"{prob_percent}%"))
 
-    if character_tags_with_probs:
-        print(f"版権キャラクター: {character_tags_with_probs}の可能性があります")
-    else:
-        print("版権キャラクターの可能性が低いと思われます")
+    general_tags_with_probs = {}
+    for i, p in enumerate(prob[4:4+len(general_tags)]):
+        if p >= thresh:
+            tag_index = i - 4
+            if tag_index < len(general_tags):
+                tag_name = general_tags[tag_index]
+                general_tags_with_probs[tag_name] = p
+    tags = sorted(general_tags_with_probs.items(), key = lambda x : x[1], reverse=True)
+
+    return {
+        'has_nsfw': max_nsfw_score > max_sfw_score,
+        'tag_confidences': tag_confidences,
+        'character_tags_with_probs': character_tags_with_probs,
+        'tags': tags,
+    }
 
 def load_model(MODEL_ID = "SmilingWolf/wd-vit-tagger-v3"):
     global ort_sess, rating_tags, character_tags, general_tags

main.py

@@ -37,5 +37,5 @@ async def check_image(file: UploadFile = Form(...)):
 
     result = fn.check(Image.open(file_stream))
 
-    return {"has_nsfw": result}
+    return result