Update

app.py

@@ -31,7 +31,7 @@ def preprocess_image(image):
     image = image.astype(np.float32)
     return image
 
-def process_image(image_path, input_name, ort_sess, rating_tags, character_tags, general_tags):
+def process_image(image, ort_sess, input_name, rating_tags, character_tags, general_tags):
     thresh = 0.35
     try:
         image = Image.open(image_path)
@@ -121,23 +121,24 @@ class webui:
         print("Hugging Faceからモデルをダウンロード中")
         onnx_path = hf_hub_download(model_id, "model.onnx")
         csv_path = hf_hub_download(model_id, "selected_tags.csv")
+        ort_sess = ort.InferenceSession(onnx_path)
 
         print("ONNXモデルを実行中")
         print(f"ONNXモデルのパス: {onnx_path}")
 
-        ort_sess = ort.InferenceSession(onnx_path)
+        image = Image.open(image_path)
+        image = image.convert("RGB") if image.mode != "RGB" else image
+        image = preprocess_image(image)
 
         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"]
+        general_tags = [row[1] for row in rows 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)
+        NSFW_flag, IP_flag, tag_text = process_image(image, ort_sess, ort_sess.get_inputs()[0].name, rating_tags, character_tags, general_tags)
         return NSFW_flag, IP_flag, tag_text