Upload app.py

app.py

@@ -208,29 +208,6 @@ def prediction_to_tag(prediction, tag_dict, class_num):
     
     return general, character, artist, date, rating
 
-def prediction_to_retrieval(prediction, tag_dict, class_num, top_k):
-    prediction = prediction.view(class_num)
-    predicted_ids = (prediction>=0.005).nonzero(as_tuple=True)[0].cpu().numpy() + 1
-
-    artist = {}
-    date = {}
-    
-    for tag, value in tag_dict.items():
-        if value[2] in predicted_ids:
-            tag_value = round(prediction[value[2] - 1].item(), 6)
-            if value[1] == "artist":
-                artist[tag] = tag_value
-            elif value[1] == "date":
-                date[tag] = tag_value
-
-    artist = dict(sorted(artist.items(), key=lambda item: item[1], reverse=True))
-    artist = dict(list(artist.items())[:top_k])
-
-    if date:
-        date = {max(date, key=date.get): date[max(date, key=date.get)]}
-    
-    return artist, date
-
 def process_image(image):
     try:
         image = image.convert('RGBA')
@@ -283,7 +260,7 @@ def process_image(image):
     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]
 
-    tags_str = ", ".join(tags_list).replace("(", r"\(").replace(")", r"\)")
+    tags_str = ", ".join(character_tags_list + tags_list).replace("(", r"\(").replace(")", r"\)")
 
     return tags_str, artist_tags, character_tags, general_tags, rating, date