164 lines
3.8 KiB
Python
164 lines
3.8 KiB
Python
from .server import MessageHolder
|
|
|
|
from comfy.sd1_clip import token_weights
|
|
|
|
|
|
class AlwaysEqualProxy(str):
|
|
def __eq__(self, _):
|
|
return True
|
|
|
|
def __ne__(self, _):
|
|
return False
|
|
|
|
|
|
any = AlwaysEqualProxy("*")
|
|
|
|
|
|
class DanbooruPromptQuiz:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"input": ("STRING", {"default": "", "multiline": True}),
|
|
"answer": (
|
|
"STRING",
|
|
{"forceInput": True},
|
|
),
|
|
},
|
|
"optional": {
|
|
"any": (any, {}),
|
|
},
|
|
"hidden": {
|
|
"id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
return inputs
|
|
|
|
RETURN_TYPES = (
|
|
"FLOAT",
|
|
"FLOAT",
|
|
"FLOAT",
|
|
"FLOAT",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
)
|
|
RETURN_NAMES = (
|
|
"accuracy",
|
|
"precision",
|
|
"recall",
|
|
"f1-measure",
|
|
"correct tags (TP)",
|
|
"wrong tags (FP)",
|
|
"missing tags (FN)",
|
|
"input",
|
|
"answer",
|
|
)
|
|
FUNCTION = "check"
|
|
|
|
def check(self, id, input, answer, any):
|
|
message = MessageHolder.waitForMessage(id)
|
|
(
|
|
accuracy,
|
|
precision,
|
|
recall,
|
|
f1,
|
|
correct_tags,
|
|
wrong_tags,
|
|
missing_tags,
|
|
) = self._check(message, answer)
|
|
return (
|
|
accuracy,
|
|
precision,
|
|
recall,
|
|
f1,
|
|
correct_tags,
|
|
wrong_tags,
|
|
missing_tags,
|
|
message,
|
|
answer,
|
|
)
|
|
|
|
def _check(self, input, answer):
|
|
input_tags = []
|
|
for t, _ in token_weights(input, 1.0):
|
|
for tag in t.split(","):
|
|
tag = tag.strip()
|
|
if tag:
|
|
input_tags.append(tag)
|
|
|
|
answer_tags = []
|
|
for t, _ in token_weights(answer, 1.0):
|
|
for tag in t.split(","):
|
|
tag = tag.strip()
|
|
if tag:
|
|
answer_tags.append(tag)
|
|
|
|
correct_tags = []
|
|
wrong_tags = []
|
|
missing_tags = []
|
|
for t in input_tags:
|
|
if t in answer_tags and t not in correct_tags:
|
|
correct_tags.append(t)
|
|
elif t not in answer_tags and t not in wrong_tags:
|
|
wrong_tags.append(t)
|
|
for t in answer_tags:
|
|
if t not in input_tags and t not in missing_tags:
|
|
missing_tags.append(t)
|
|
|
|
tp = len(correct_tags)
|
|
fp = len(wrong_tags)
|
|
fn = len(missing_tags)
|
|
accuracy = tp / (tp + fp + fn)
|
|
precision = tp / (tp + fp)
|
|
recall = tp / (tp + fn)
|
|
f1 = 2 * precision * recall / (precision + recall)
|
|
|
|
return (
|
|
accuracy,
|
|
precision,
|
|
recall,
|
|
f1,
|
|
", ".join(correct_tags),
|
|
", ".join(wrong_tags),
|
|
", ".join(missing_tags),
|
|
)
|
|
|
|
|
|
class DanbooruPromptComparison(DanbooruPromptQuiz):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"input": ("STRING", {"default": "", "multiline": True}),
|
|
"answer": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
return inputs
|
|
|
|
def check(self, input, answer):
|
|
(
|
|
accuracy,
|
|
precision,
|
|
recall,
|
|
f1,
|
|
correct_tags,
|
|
wrong_tags,
|
|
missing_tags,
|
|
) = self._check(input, answer)
|
|
return (
|
|
accuracy,
|
|
precision,
|
|
recall,
|
|
f1,
|
|
correct_tags,
|
|
wrong_tags,
|
|
missing_tags,
|
|
input,
|
|
answer,
|
|
)
|