From: Project Hentai AI Date: Wed, 2 Feb 2022 12:29:50 +0000 (+0000) Subject: Code, lists and demo dataset added X-Git-Url: https://git.hentai-ai.org/?a=commitdiff_plain;p=hentai-tagger.git%2F.git Code, lists and demo dataset added --- diff --git a/categories.txt b/categories.txt new file mode 100644 index 0000000..e16c592 --- /dev/null +++ b/categories.txt @@ -0,0 +1,3 @@ +hentai +ecchi +lewd \ No newline at end of file diff --git a/demo_input_dir/20211105_211129.jpg b/demo_input_dir/20211105_211129.jpg new file mode 100644 index 0000000..7e2f5b0 Binary files /dev/null and b/demo_input_dir/20211105_211129.jpg differ diff --git a/demo_input_dir/3cOD2j5.jpg b/demo_input_dir/3cOD2j5.jpg new file mode 100644 index 0000000..b1abf68 Binary files /dev/null and b/demo_input_dir/3cOD2j5.jpg differ diff --git a/demo_input_dir/410acd8a76deaa121a3ce9b1c7c45786.jpg b/demo_input_dir/410acd8a76deaa121a3ce9b1c7c45786.jpg new file mode 100644 index 0000000..6c1fdb5 Binary files /dev/null and b/demo_input_dir/410acd8a76deaa121a3ce9b1c7c45786.jpg differ diff --git a/demo_input_dir/IMG_5688.jpg b/demo_input_dir/IMG_5688.jpg new file mode 100644 index 0000000..4cd16d7 Binary files /dev/null and b/demo_input_dir/IMG_5688.jpg differ diff --git a/demo_input_dir/cf550cbdb25084e9728a3d38e430d0a3.jpg b/demo_input_dir/cf550cbdb25084e9728a3d38e430d0a3.jpg new file mode 100644 index 0000000..e0dae7c Binary files /dev/null and b/demo_input_dir/cf550cbdb25084e9728a3d38e430d0a3.jpg differ diff --git a/demo_input_dir/dfb4f842aca3409938305ecd5e24d772.png b/demo_input_dir/dfb4f842aca3409938305ecd5e24d772.png new file mode 100644 index 0000000..b21ffbe Binary files /dev/null and b/demo_input_dir/dfb4f842aca3409938305ecd5e24d772.png differ diff --git a/demo_input_dir/ougu2m62lkd81.jpg b/demo_input_dir/ougu2m62lkd81.jpg new file mode 100644 index 0000000..4a49f04 Binary files /dev/null and b/demo_input_dir/ougu2m62lkd81.jpg differ diff --git a/demo_input_dir/sample_0403990a44aad89becce254964ea89c8.jpg b/demo_input_dir/sample_0403990a44aad89becce254964ea89c8.jpg new file mode 100644 index 0000000..754ef40 Binary files /dev/null and b/demo_input_dir/sample_0403990a44aad89becce254964ea89c8.jpg differ diff --git a/demo_input_dir/sample_174aad2b4c9db715d7373832229ee52e.jpg b/demo_input_dir/sample_174aad2b4c9db715d7373832229ee52e.jpg new file mode 100644 index 0000000..bbcc673 Binary files /dev/null and b/demo_input_dir/sample_174aad2b4c9db715d7373832229ee52e.jpg differ diff --git a/demo_input_dir/sample_f4a72d77d72e7c7329ac4af45eb3a78f.jpg b/demo_input_dir/sample_f4a72d77d72e7c7329ac4af45eb3a78f.jpg new file mode 100644 index 0000000..eec3d54 Binary files /dev/null and b/demo_input_dir/sample_f4a72d77d72e7c7329ac4af45eb3a78f.jpg differ diff --git a/hentai-tagger.py b/hentai-tagger.py new file mode 100644 index 0000000..75bdf8a --- /dev/null +++ b/hentai-tagger.py @@ -0,0 +1,218 @@ +from tkinter import * +from PIL import ImageTk, Image +import argparse +import math +import os +import sys + +# Todo: + +root = Tk() +root.title('Tag App') + +parser = argparse.ArgumentParser() +parser.add_argument("-i", required=True, default="in_dir", help="input directory with images") +parser.add_argument("-o", required=False, default="", help="output directory for images") +parser.add_argument("-t", required=True, default="tags.txt", help="text file with tags") +parser.add_argument("-c", required=False, default="categories.txt", help="text file with categories") +parser.add_argument("--delete", action='store_true', help="whether to delete original files") +args = vars(parser.parse_args()) + +in_folder = args["i"] +out_folder = args["o"] if os.path.isdir(args["o"]) else args["i"] +tags_file = args["t"] +cat_file = args["c"] +delete_original = args["delete"] + +win_w = 800 +win_h = 600 + + +def resize_img(img): + img_w, img_h = img.size + if img_w > img_h: + return ImageTk.PhotoImage( + img.resize((win_w, int(win_w * img_h / img_w)), Image.ANTIALIAS)) + else: + return ImageTk.PhotoImage( + img.resize((int(win_h * img_w / img_h), win_h), Image.ANTIALIAS)) + + +class MainWindow: + + def __init__(self, main): + self.canvas = Canvas(main, width=win_w, height=win_h) + self.canvas.pack() + self.main = main + + # Load images as a list of PIL objects + self.image_list = [] + self.original_fnames = [f for f in os.listdir(in_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))] + for fname in self.original_fnames: + self.image_list.append(Image.open(os.path.join(in_folder, fname))) + print(f"loaded {len(self.image_list)} images") + + # Render first image on canvas + self.image_number = 0 + first_img = resize_img(self.image_list[self.image_number]) + self.img_label = Label(image=first_img) + self.img_label.image = first_img + self.img_label.place(anchor=NW) + + self.top_frame = Frame(main) + self.top_frame.pack(side=TOP) + + # Next button + self.next_button = Button(self.top_frame, text="Next", bg="lightblue", width=10, command=self.next) + self.next_button.pack(side=RIGHT) + + # Load tags from txt file + if os.path.isfile(tags_file): + with open(tags_file, "r") as t: + self.tags = t.read().splitlines() + self.tags.sort() + else: + sys.exit(f"{tags_file} not found") + + self.button_list = [] + self.bool_list = [] + self.frames = [] + self.buttons_per_row = 10 + self.create_tag_buttons() + self.default_color = self.button_list[0].cget("background") + + # Load categories from txt file + self.categories = [] + if os.path.isfile(cat_file): + with open(cat_file, "r") as c: + self.categories = c.read().splitlines() + + # Create an OptionMenu from categories + if len(self.categories) > 0: + self.variable = StringVar(main) + self.variable.set(self.categories[0]) # default value + self.op_menu = OptionMenu(self.top_frame, self.variable, *self.categories) + self.op_menu.pack(side=RIGHT) + + # Create an input textfield for names or custom tags + char_frame = Frame(main) + char_frame.pack(side=TOP) + self.char_label = Label(char_frame, text='Character') + self.char_label.pack(side=LEFT) + self.char_name = Text(char_frame, width=15, height=1) + self.char_name.pack(side=LEFT) + + artist_frame = Frame(main) + artist_frame.pack(side=TOP) + self.artist_label = Label(artist_frame, text='Artist') + self.artist_label.pack(side=LEFT) + self.artist_name = Text(artist_frame, width=15, height=1) + self.artist_name.pack(side=LEFT) + + source_frame = Frame(main) + source_frame.pack(side=TOP) + self.source_label = Label(source_frame, text='Source') + self.source_label.pack(side=LEFT) + self.source_name = Text(source_frame, width=15, height=1) + self.source_name.pack(side=LEFT) + + + # Create an input textfield for new tags + new_tags_frame = Frame(main) + new_tags_frame.pack(side=TOP) + self.new_tags_box = Text(new_tags_frame, width=15, height=1) + self.new_tags_box.pack(side=LEFT) + self.new_tags_button = Button(new_tags_frame, text="Add Tag", command=self.add_new_tag) + self.new_tags_button.pack(side=LEFT) + + # Button function for toggle + def switch(self, button_id): + if self.bool_list[button_id]: + self.button_list[button_id].config(bg=self.default_color) + self.bool_list[button_id] = False + else: + self.button_list[button_id].config(bg='gray55') + self.bool_list[button_id] = True + + # Function for renaming image file + def rename(self): + counter = 1 + tag_string = "" + for x in range(len(self.tags)): + if self.bool_list[x]: + tag_string += self.tags[x] + "_" + tag_string = tag_string[:-1] + if len(self.categories) > 0: + tag_string = self.variable.get() + "_" + tag_string if tag_string != "" else self.variable.get() + + name_list = [] + name_list.append(self.char_name.get("1.0", 'end-1c').lower().replace(" ", "_")) + name_list.append(self.artist_name.get("1.0", 'end-1c').lower().replace(" ", "_")) + name_list.append(self.source_name.get("1.0", 'end-1c').lower().replace(" ", "_")) + name_list = [x for x in name_list if x != ""] + tag_string = tag_string + "-" + "_".join(name_list) if len(name_list) > 0 else tag_string + + new_files = os.listdir(out_folder) + new_fname = tag_string + "-" + str(counter) + ".png" + while new_fname in new_files: + counter += 1 + new_fname = '-'.join(new_fname.split('.')[0].split('-')[:-1]) + "-" + str(counter) + ".png" + self.image_list[self.image_number].save(os.path.join(out_folder, new_fname)) + if delete_original: + os.remove(os.path.join(in_folder, self.original_fnames[self.image_number])) + + # Function for changing image + def next(self): + self.rename() + + self.image_number += 1 + if self.image_number == len(self.image_list): + print("exiting...") + root.destroy() + sys.exit(0) + next_img = resize_img(self.image_list[self.image_number]) + self.img_label.config(image=next_img) + self.img_label.image = next_img + + # Reset toggle buttons and input textfield + for x in range(len(self.tags)): + self.button_list[x].config(bg=self.default_color) + self.bool_list[x] = False + self.char_name.delete("1.0", END) + self.artist_name.delete("1.0", END) + self.source_name.delete("1.0", END) + + def add_new_tag(self): + self.tags.append(self.new_tags_box.get("1.0", 'end-1c').lower()) + self.new_tags_box.delete('1.0', END) + self.tags.sort() + self.create_tag_buttons() + with open(tags_file, "w") as f: + f.write('\n'.join(self.tags) + '\n') + + def create_tag_buttons(self): + # Remove previous buttons/frames + for b in self.button_list: + b.pack_forget() + for f in self.frames: + f.pack_forget() + self.button_list = [] + self.bool_list = [] + self.frames = [] + + # Create a toggle button for each tag + for f in range(int(math.ceil(len(self.tags) / self.buttons_per_row))): + tmp_f = Frame(self.main) + tmp_f.pack(side=BOTTOM) + self.frames.append(tmp_f) + self.frames.reverse() + + for i in range(len(self.tags)): + b = Button(self.frames[int(i / self.buttons_per_row)], text=self.tags[i], command=lambda i=i: self.switch(i)) + b.pack(side=LEFT) + self.button_list.append(b) + self.bool_list.append(False) + + +MainWindow(root) +root.mainloop() diff --git a/tags.txt b/tags.txt new file mode 100644 index 0000000..385e925 --- /dev/null +++ b/tags.txt @@ -0,0 +1,30 @@ +blindfold +bondage +buttplug +choker +collar +darkskin +demongirl +elf +glasses +inmon +juice +kemonomimi +keyhole +kneehighs +leash +leggings +legslitdress +maid +masturbation +monstergirl +onsen +selfie +swimsuit +tanlines +thighhighs +tomboy +toy +underboob +virginkiller +yuri