Upload app.py

app.py

@@ -151,6 +151,9 @@ with open(os.path.join(repo_dir, 'character_threshold.json'), 'r', encoding='utf
 with open(os.path.join(repo_dir, 'general_threshold.json'), 'r', encoding='utf-8') as f:
     general_thresholds = json.load(f)
 
+with open(os.path.join(repo_dir, 'character_feature.json'), 'r', encoding='utf-8') as f:
+    character_features = json.load(f)
+
 model_map = MultiheadAttentionPoolingHead(2048)
 model_map.load_state_dict(torch.load(os.path.join(repo_dir, "map_head.pth"), map_location=device, weights_only=True))
 model_map.to(device).to(dtype).eval()
@@ -253,6 +256,13 @@ def process_image(image):
         character_ = prediction_to_tag(character_prediction, character_dict, character_class)
         character_tags = character_[1]
 
+        remove_list = []
+        for tag in character_tags:
+            if tag in implications_list:
+                remove_list.extend([implication for implication in implications_list[tag]])
+
+        character_tags_list = [tag for tag in character_tags if tag not in remove_list]
+
         artist_prediction = mlp_artist(embedding)
         artist_ = prediction_to_tag(artist_prediction, artist_dict, artist_class)
         artist_tags = artist_[2]
@@ -266,6 +276,10 @@ def process_image(image):
         if tag in implications_list:
             for implication in implications_list[tag]:
                 remove_list.append(implication)
+    for char_tag in character_tags_list:
+        if char_tag in character_features:
+            for character_feature in character_features[char_tag]:
+                remove_list.append(character_feature)
     tags_list = [tag for tag in tags_list if tag not in remove_list]
     tags_list = [tag.replace("_", " ") if tag not in kaomojis else tag for tag in tags_list]