Skip to content

Commit

Permalink
Important Auto-Tag Model Patch
Browse files Browse the repository at this point in the history
  • Loading branch information
@x-CK-x
x-CK-x authored Jul 21, 2023
1 parent f0eb354 commit fbc218c
Showing 1 changed file with 19 additions and 15 deletions.
34 changes: 19 additions & 15 deletions autotag.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,9 +140,12 @@ def run_batch(self, path_imgs, single_image, all_tags_ever_dict):
if single_image:
found_tags = self.label_names[self.label_names["probs"] > float(0)][["name", "probs"]]
found_tags = found_tags.values.tolist()

# help.verbose_print(f"found_tags:\t{found_tags}")

# remove tags not in csv or not a valid category type
found_tags = [tag for tag in found_tags if
(self.categories_map[all_tags_ever_dict[tag[0]]] in self.valid_categories) and (tag[0] in all_tags_ever_dict)]
found_tags = [tag for tag in found_tags if (tag[0] in all_tags_ever_dict) and
(self.categories_map[all_tags_ever_dict[tag[0]][0]] in self.valid_categories)]

# set predictions for the UI
for element in found_tags:
Expand All @@ -158,15 +161,15 @@ def run_batch(self, path_imgs, single_image, all_tags_ever_dict):
# convert to list
found_tags = found_tags.values.tolist()
# remove tags not in csv or not a valid category type
found_tags = [tag for tag in found_tags if
(self.categories_map[all_tags_ever_dict[tag[0]]] in self.valid_categories) and (tag[0] in all_tags_ever_dict)]
found_tags = [tag for tag in found_tags if (tag[0] in all_tags_ever_dict) and
(self.categories_map[all_tags_ever_dict[tag[0]][0]] in self.valid_categories)]

# set predictions for the UI
for element in found_tags:
self.combined_tags[element[0]] = [element[1], f"{image_name}.{image_ext}"] # tag -> [probability, name w/ extension]

# sort by category
sorted_list = sorted(found_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x[0]]]])
sorted_list = sorted(found_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x[0]][0]]])
sorted_list = [pair[0] for pair in sorted_list] # get just the tags

# Load image tags from file and sort them into the same kind of this sorted by category
Expand All @@ -178,17 +181,17 @@ def run_batch(self, path_imgs, single_image, all_tags_ever_dict):
f = open(temp_src_tags_path, 'w')
f.close()
# sort by category
sorted_existing_tags_list = sorted(existing_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]])
sorted_existing_tags_list = sorted(existing_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]])

# remove duplicate tag/s in generated tag list
sorted_list1_set = set(sorted_existing_tags_list)
filtered_sorted_list2 = [tag for tag in sorted_list if tag not in sorted_list1_set]
filtered_sorted_list2 = sorted(filtered_sorted_list2, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]])
filtered_sorted_list2 = sorted(filtered_sorted_list2, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]])

merged_list = []
if self.write_tag_opts_dropdown == 'Merge':
# merge the two sorted lists by category
merged_list = list(heapq.merge(sorted_existing_tags_list, filtered_sorted_list2, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]]))
merged_list = list(heapq.merge(sorted_existing_tags_list, filtered_sorted_list2, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]]))
elif self.write_tag_opts_dropdown == 'Pre-pend':
# pre-pend the generated list to the existing one
merged_list = filtered_sorted_list2 + sorted_existing_tags_list
Expand Down Expand Up @@ -239,7 +242,7 @@ def run_batch(self, path_imgs, single_image, all_tags_ever_dict):
copy.deepcopy(meta_csv_dict),
copy.deepcopy(rating_csv_dict),
copy.deepcopy(tags_csv_dict),
self.categories_map[all_tags_ever_dict[tag]], tag, "+", 1)
self.categories_map[all_tags_ever_dict[tag][0]], tag, "+", 1)
# persist changes to csv dictionary files
help.write_tags_to_csv(artist_csv_dict, os.path.join(self.tag_folder, "artist.csv"))
help.write_tags_to_csv(character_csv_dict, os.path.join(self.tag_folder, "character.csv"))
Expand Down Expand Up @@ -382,10 +385,11 @@ def save_tags(self, single_image, any_selected_tags, all_tags_ever_dict):
raise ValueError("batch use tag operation NOT set")

# remove tags not in csv or not a valid category type
any_selected_tags = [tag for tag in any_selected_tags if (self.categories_map[all_tags_ever_dict[tag]] in self.valid_categories) and (tag in all_tags_ever_dict)]
any_selected_tags = [tag for tag in any_selected_tags if (tag in all_tags_ever_dict) and
(self.categories_map[all_tags_ever_dict[tag][0]] in self.valid_categories)]

# sort by category
sorted_list = sorted(any_selected_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]])
sorted_list = sorted(any_selected_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]])

# Load image tags from file and sort them into the same kind of this sorted by category
existing_tags = []
Expand All @@ -396,20 +400,20 @@ def save_tags(self, single_image, any_selected_tags, all_tags_ever_dict):
f = open(temp_src_tags_path, 'w')
f.close()
# sort by category
sorted_existing_tags_list = sorted(existing_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]])
sorted_existing_tags_list = sorted(existing_tags, key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]])

# remove duplicate tag/s in generated tag list
sorted_list1_set = set(sorted_existing_tags_list)
filtered_sorted_list2 = [tag for tag in sorted_list if tag not in sorted_list1_set]

filtered_sorted_list2 = sorted(filtered_sorted_list2,
key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]])
key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]])

merged_list = []
if self.write_tag_opts_dropdown == 'Merge':
# merge the two sorted lists by category
merged_list = list(heapq.merge(sorted_existing_tags_list, filtered_sorted_list2,
key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x]]]))
key=lambda x: self.valid_categories[self.categories_map[all_tags_ever_dict[x][0]]]))
elif self.write_tag_opts_dropdown == 'Pre-pend':
# pre-pend the generated list to the existing one
merged_list = filtered_sorted_list2 + sorted_existing_tags_list
Expand Down Expand Up @@ -461,7 +465,7 @@ def save_tags(self, single_image, any_selected_tags, all_tags_ever_dict):
copy.deepcopy(meta_csv_dict),
copy.deepcopy(rating_csv_dict),
copy.deepcopy(tags_csv_dict),
self.categories_map[all_tags_ever_dict[tag]], tag, "+", 1)
self.categories_map[all_tags_ever_dict[tag][0]], tag, "+", 1)
# persist changes to csv dictionary files
help.write_tags_to_csv(artist_csv_dict, os.path.join(self.tag_folder, "artist.csv"))
help.write_tags_to_csv(character_csv_dict, os.path.join(self.tag_folder, "character.csv"))
Expand Down

0 comments on commit fbc218c

Please sign in to comment.