import tkinter as tk from tkinter import ttk, Menu import os import subprocess from PIL import Image, ImageTk, ImageDraw import tkinter.filedialog as fd import json import sys import os import sys from torchvision import transforms from torchvision.transforms.functional import InterpolationMode import torch import subprocess import numpy as np import requests import random import customtkinter as ctk from customtkinter import ThemeManager from clip_segmentation import ClipSeg #main class ctk.set_appearance_mode("dark") ctk.set_default_color_theme("blue") class BatchMaskWindow(ctk.CTkToplevel): def __init__(self, parent, path, *args, **kwargs): ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) self.parent = parent self.title("Batch process masks") self.geometry("320x310") self.resizable(False, False) self.wait_visibility() self.grab_set() self.focus_set() self.mode_var = tk.StringVar(self, "Create if absent") self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing"] self.frame = ctk.CTkFrame(self, width=600, height=300) self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) self.path_label.grid(row=0, column=0, sticky="w",padx=5, pady=5) self.path_entry = ctk.CTkEntry(self.frame, width=150) self.path_entry.insert(0, path) self.path_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) self.path_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) self.prompt_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) self.prompt_entry = ctk.CTkEntry(self.frame, width=200) self.prompt_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) self.mode_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) self.mode_dropdown.grid(row=2, column=1, sticky="w", padx=5, pady=5) self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) self.threshold_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") self.threshold_entry.insert(0, "0.3") self.threshold_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) self.smooth_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") self.smooth_entry.insert(0, 5) self.smooth_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) self.expand_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") self.expand_entry.insert(0, 10) self.expand_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) self.progress_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) self.progress.grid(row=6, column=1, sticky="w", padx=5, pady=5) self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) self.create_masks_button.grid(row=7, column=0, columnspan=2, sticky="w", padx=5, pady=5) self.frame.pack(fill="both", expand=True) def browse_for_path(self, entry_box): # get the path from the user path = fd.askdirectory() # set the path to the entry box # delete entry box text entry_box.focus_set() entry_box.delete(0, tk.END) entry_box.insert(0, path) self.focus_set() def set_progress(self, value, max_value): progress = value / max_value self.progress.set(progress) self.progress_label.configure(text="{0}/{1}".format(value, max_value)) self.progress.update() def create_masks(self): self.parent.load_clip_seg_model() mode = { "Replace all masks": "replace", "Create if absent": "fill", "Add to existing": "add", "Subtract from existing": "subtract" }[self.mode_var.get()] self.parent.clip_seg.mask_folder( sample_dir=self.path_entry.get(), prompts=[self.prompt_entry.get()], mode=mode, threshold=float(self.threshold_entry.get()), smooth_pixels=int(self.smooth_entry.get()), expand_pixels=int(self.expand_entry.get()), progress_callback=self.set_progress, ) self.parent.load_image() def _check_file_type(f: str) -> bool: return f.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', ".bmp", ".tiff")) class ImageBrowser(ctk.CTkToplevel): def __init__(self,mainProcess=None): super().__init__() if not os.path.exists("scripts/BLIP"): print("Getting BLIP from GitHub.") subprocess.run(["git", "clone", "https://github.com/salesforce/BLIP", "scripts/BLIP"]) #if not os.path.exists("scripts/CLIP"): # print("Getting CLIP from GitHub.") # subprocess.run(["git", "clone", "https://github.com/pharmapsychotic/clip-interrogator.git', 'scripts/CLIP"]) blip_path = "scripts/BLIP" sys.path.append(blip_path) #clip_path = "scripts/CLIP" #sys.path.append(clip_path) self.mainProcess = mainProcess self.captioner_folder = os.path.dirname(os.path.realpath(__file__)) self.clip_seg = None self.PILimage = None self.PILmask = None self.mask_draw_x = 0 self.mask_draw_y = 0 self.mask_draw_radius = 20 #self = master #self.overrideredirect(True) #self.title_bar = TitleBar(self) #self.title_bar.pack(side="top", fill="x") #make not user resizable self.title("Caption Buddy") #self.resizable(False, False) self.geometry("720x820") self.top_frame = ctk.CTkFrame(self,fg_color='transparent') self.top_frame.pack(side="top", fill="x",expand=False) self.top_subframe = ctk.CTkFrame(self.top_frame,fg_color='transparent') self.top_subframe.pack(side="bottom", fill="x",pady=10) self.top_subframe.grid_columnconfigure(0, weight=1) self.top_subframe.grid_columnconfigure(1, weight=1) self.tip_frame = ctk.CTkFrame(self,fg_color='transparent') self.tip_frame.pack(side="top") self.dark_mode_var = "#202020" #self.dark_purple_mode_var = "#1B0F1B" self.dark_mode_title_var = "#286aff" self.dark_mode_button_pressed_var = "#BB91B6" self.dark_mode_button_var = "#8ea0e1" self.dark_mode_text_var = "#c6c7c8" #self.configure(bg_color=self.dark_mode_var) self.canvas = ctk.CTkLabel(self,text='', width=600, height=600) #self.canvas.configure(bg_color=self.dark_mode_var) #create temporary image for canvas self.canvas.pack() self.cur_img_index = 0 self.image_count = 0 #make a frame with a grid under the canvas self.frame = ctk.CTkFrame(self) #grid self.frame.grid_columnconfigure(0, weight=1) self.frame.grid_columnconfigure(1, weight=100) self.frame.grid_columnconfigure(2, weight=1) self.frame.grid_rowconfigure(0, weight=1) #show the frame self.frame.pack(side="bottom", fill="x") #bottom frame self.bottom_frame = ctk.CTkFrame(self) #make grid self.bottom_frame.grid_columnconfigure(0, weight=0) self.bottom_frame.grid_columnconfigure(1, weight=2) self.bottom_frame.grid_columnconfigure(2, weight=0) self.bottom_frame.grid_columnconfigure(3, weight=2) self.bottom_frame.grid_columnconfigure(4, weight=0) self.bottom_frame.grid_columnconfigure(5, weight=2) self.bottom_frame.grid_rowconfigure(0, weight=1) #show the frame self.bottom_frame.pack(side="bottom", fill="x") self.image_index = 0 self.image_list = [] self.caption = '' self.caption_file = '' self.caption_file_path = '' self.caption_file_name = '' self.caption_file_ext = '' self.caption_file_name_no_ext = '' self.output_format='text' #check if bad_files.txt exists if os.path.exists("bad_files.txt"): #delete it os.remove("bad_files.txt") self.use_blip = True self.debug = False self.create_widgets() self.load_blip_model() self.load_options() #self.open_folder() self.canvas.focus_force() self.canvas.bind("", self.next_image) self.canvas.bind("", self.prev_image) #on close window self.protocol("WM_DELETE_WINDOW", self.on_closing) def on_closing(self): #self.save_options() self.mainProcess.deiconify() self.destroy() def create_widgets(self): self.output_folder = '' # add a checkbox to toggle auto generate caption self.auto_generate_caption = tk.BooleanVar(self.top_subframe) self.auto_generate_caption.set(True) self.auto_generate_caption_checkbox = ctk.CTkCheckBox(self.top_subframe, text="Auto Generate Caption", variable=self.auto_generate_caption,width=50) self.auto_generate_caption_checkbox.pack(side="left", fill="x", expand=True, padx=10) # add a checkbox to skip auto generating captions if they already exist self.auto_generate_caption_text_override = tk.BooleanVar(self.top_subframe) self.auto_generate_caption_text_override.set(False) self.auto_generate_caption_checkbox_text_override = ctk.CTkCheckBox(self.top_subframe, text="Skip Auto Generate If Text Caption Exists", variable=self.auto_generate_caption_text_override,width=50) self.auto_generate_caption_checkbox_text_override.pack(side="left", fill="x", expand=True, padx=10) # add a checkbox to enable mask editing self.enable_mask_editing = tk.BooleanVar(self.top_subframe) self.enable_mask_editing.set(False) self.enable_mask_editing_checkbox = ctk.CTkCheckBox(self.top_subframe, text="Enable Mask Editing", variable=self.enable_mask_editing, width=50) self.enable_mask_editing_checkbox.pack(side="left", fill="x", expand=True, padx=10) self.open_button = ctk.CTkButton(self.top_frame,text="Load Folder",fg_color=("gray75", "gray25"), command=self.open_folder,width=50) #self.open_button.grid(row=0, column=1) self.open_button.pack(side="left", fill="x",expand=True,padx=10) #add a batch folder button self.batch_folder_caption_button = ctk.CTkButton(self.top_frame, text="Batch Folder Caption", fg_color=("gray75", "gray25"), command=self.batch_folder_caption, width=50) self.batch_folder_caption_button.pack(side="left", fill="x", expand=True, padx=10) self.batch_folder_mask_button = ctk.CTkButton(self.top_frame, text="Batch Folder Mask", fg_color=("gray75", "gray25"), command=self.batch_folder_mask, width=50) self.batch_folder_mask_button.pack(side="left", fill="x", expand=True, padx=10) #add an options button to the same row as the open button self.options_button = ctk.CTkButton(self.top_frame, text="Options",fg_color=("gray75", "gray25"), command=self.open_options,width=50) self.options_button.pack(side="left", fill="x",expand=True,padx=10) #add generate caption button self.generate_caption_button = ctk.CTkButton(self.top_frame, text="Generate Caption",fg_color=("gray75", "gray25"), command=self.generate_caption,width=50) self.generate_caption_button.pack(side="left", fill="x",expand=True,padx=10) #add a label for tips under the buttons self.tips_label = ctk.CTkLabel(self.tip_frame, text="Use Alt with left and right arrow keys to navigate images, enter to save the caption.") self.tips_label.pack(side="top") #add image count label self.image_count_label = ctk.CTkLabel(self.tip_frame, text=f"Image {self.cur_img_index} of {self.image_count}") self.image_count_label.pack(side="top") self.image_label = ctk.CTkLabel(self.canvas,text='',width=100,height=100) self.image_label.grid(row=0, column=0, sticky="nsew") #self.image_label.bind("", self.click_canvas) self.image_label.bind("", self.draw_mask) self.image_label.bind("", self.draw_mask) self.image_label.bind("", self.draw_mask) self.image_label.bind("", self.draw_mask_radius) #self.image_label.pack(side="top") #previous button self.prev_button = ctk.CTkButton(self.frame,text="Previous", command= lambda event=None: self.prev_image(event),width=50) #grid self.prev_button.grid(row=1, column=0, sticky="w",padx=5,pady=10) #self.prev_button.pack(side="left") #self.prev_button.bind("", self.prev_image) self.caption_entry = ctk.CTkEntry(self.frame) #grid self.caption_entry.grid(row=1, column=1, rowspan=3, sticky="nsew",pady=10) #bind to enter key self.caption_entry.bind("", self.save) self.canvas.bind("", self.save) self.caption_entry.bind("", self.next_image) self.caption_entry.bind("", self.prev_image) self.caption_entry.bind("", self.delete_word) #next button self.next_button = ctk.CTkButton(self.frame,text='Next', command= lambda event=None: self.next_image(event),width=50) #self.next_button["text"] = "Next" #grid self.next_button.grid(row=1, column=2, sticky="e",padx=5,pady=10) #add two entry boxes and labels in the style of :replace _ with _ #create replace string variable self.replace_label = ctk.CTkLabel(self.bottom_frame, text="Replace:") self.replace_label.grid(row=0, column=0, sticky="w",padx=5) self.replace_entry = ctk.CTkEntry(self.bottom_frame, ) self.replace_entry.grid(row=0, column=1, sticky="nsew",padx=5) self.replace_entry.bind("", self.save) #self.replace_entry.bind("", self.replace) #with label #create with string variable self.with_label = ctk.CTkLabel(self.bottom_frame, text="With:") self.with_label.grid(row=0, column=2, sticky="w",padx=5) self.with_entry = ctk.CTkEntry(self.bottom_frame, ) self.with_entry.grid(row=0, column=3, sticky="nswe",padx=5) self.with_entry.bind("", self.save) #add another entry with label, add suffix #create prefix string var self.prefix_label = ctk.CTkLabel(self.bottom_frame, text="Add to start:") self.prefix_label.grid(row=0, column=4, sticky="w",padx=5) self.prefix_entry = ctk.CTkEntry(self.bottom_frame, ) self.prefix_entry.grid(row=0, column=5, sticky="nsew",padx=5) self.prefix_entry.bind("", self.save) #create suffix string var self.suffix_label = ctk.CTkLabel(self.bottom_frame, text="Add to end:") self.suffix_label.grid(row=0, column=6, sticky="w",padx=5) self.suffix_entry = ctk.CTkEntry(self.bottom_frame, ) self.suffix_entry.grid(row=0, column=7, sticky="nsew",padx=5) self.suffix_entry.bind("", self.save) self.all_entries = [self.replace_entry, self.with_entry, self.suffix_entry, self.caption_entry, self.prefix_entry] #bind right click menu to all entries for entry in self.all_entries: entry.bind("", self.create_right_click_menu) def batch_folder_caption(self): #show imgs in folder askdirectory #ask user if to batch current folder or select folder #if bad_files.txt exists, delete it self.bad_files = [] if os.path.exists('bad_files.txt'): os.remove('bad_files.txt') try: #check if self.folder is set self.folder except AttributeError: self.folder = '' if self.folder == '': self.folder = fd.askdirectory(title="Select Folder to Batch Process", initialdir=os.getcwd()) batch_input_dir = self.folder else: ask = tk.messagebox.askquestion("Batch Folder", "Batch current folder?") if ask == 'yes': batch_input_dir = self.folder else: batch_input_dir = fd.askdirectory(title="Select Folder to Batch Process", initialdir=os.getcwd()) ask2 = tk.messagebox.askquestion("Batch Folder", "Save output to same directory?") if ask2 == 'yes': batch_output_dir = batch_input_dir else: batch_output_dir = fd.askdirectory(title="Select Folder to Save Batch Processed Images", initialdir=os.getcwd()) if batch_input_dir == '': return if batch_output_dir == '': batch_output_dir = batch_input_dir self.caption_file_name = os.path.basename(batch_input_dir) self.image_list = [] for file in os.listdir(batch_input_dir): if _check_file_type(file) and not file.endswith('-masklabel.png'): self.image_list.append(os.path.join(batch_input_dir, file)) self.image_index = 0 #use progress bar class #pba = tk.Tk() #pba.title("Batch Processing") #remove icon #pba.wm_attributes('-toolwindow','True') pb = ProgressbarWithCancel(max=len(self.image_list)) #pb.set_max(len(self.image_list)) pb.set_progress(0) #if batch_output_dir doesn't exist, create it if not os.path.exists(batch_output_dir): os.makedirs(batch_output_dir) for i in range(len(self.image_list)): radnom_chance = random.randint(0,25) if radnom_chance == 0: pb.set_random_label() if pb.is_cancelled(): pb.destroy() return self.image_index = i #get float value of progress between 0 and 1 according to the image index and the total number of images progress = i / len(self.image_list) pb.set_progress(progress) self.update() try: img = Image.open(self.image_list[i]).convert("RGB") except: self.bad_files.append(self.image_list[i]) #skip file continue tensor = transforms.Compose([ transforms.Resize((self.blipSize, self.blipSize), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) torch_image = tensor(img).unsqueeze(0).to(torch.device("cuda")) if self.nucleus_sampling: captions = self.blip_decoder.generate(torch_image, sample=True, top_p=self.q_factor) else: captions = self.blip_decoder.generate(torch_image, sample=False, num_beams=16, min_length=self.min_length, \ max_length=48, repetition_penalty=self.q_factor) caption = captions[0] self.replace = self.replace_entry.get() self.replace_with = self.with_entry.get() self.suffix_var = self.suffix_entry.get() self.prefix = self.prefix_entry.get() #prepare the caption if self.suffix_var.startswith(',') or self.suffix_var.startswith(' '): self.suffix_var = self.suffix_var else: self.suffix_var = ' ' + self.suffix_var caption = caption.replace(self.replace, self.replace_with) if self.prefix != '': if self.prefix.endswith(' '): self.prefix = self.prefix[:-1] if not self.prefix.endswith(','): self.prefix = self.prefix+',' caption = self.prefix + ' ' + caption if caption.endswith(',') or caption.endswith('.'): caption = caption[:-1] caption = caption +', ' + self.suffix_var else: caption = caption + self.suffix_var #saving the captioned image if self.output_format == 'text': #text file with same name as image imgName = os.path.basename(self.image_list[self.image_index]) imgName = imgName[:imgName.rfind('.')] caption_file = os.path.join(batch_output_dir, imgName + '.txt') with open(caption_file, 'w') as f: f.write(caption) elif self.output_format == 'filename': #duplicate image with caption as file name img.save(os.path.join(batch_output_dir, caption+'.png')) progress = i + 1 / len(self.image_list) pb.set_progress(progress) #show message box when done pb.destroy() donemsg = tk.messagebox.showinfo("Batch Folder", "Batching complete!",parent=self.master) if len(self.bad_files) > 0: bad_files_msg = tk.messagebox.showinfo("Bad Files", "Couldn't process " + str(len(self.bad_files)) + "files,\nFor a list of problematic files see bad_files.txt",parent=self.master) with open('bad_files.txt', 'w') as f: for item in self.bad_files: f.write(item + '\n') #ask user if we should load the batch output folder ask3 = tk.messagebox.askquestion("Batch Folder", "Load batch output folder?") if ask3 == 'yes': self.image_index = 0 self.open_folder(folder=batch_output_dir) #focus on donemsg #donemsg.focus_force() def generate_caption(self): #get the image tensor = transforms.Compose([ #transforms.CenterCrop(SIZE), transforms.Resize((self.blipSize, self.blipSize), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) torch_image = tensor(self.PILimage).unsqueeze(0).to(torch.device("cuda")) if self.nucleus_sampling: captions = self.blip_decoder.generate(torch_image, sample=True, top_p=self.q_factor) else: captions = self.blip_decoder.generate(torch_image, sample=False, num_beams=16, min_length=self.min_length, \ max_length=48, repetition_penalty=self.q_factor) self.caption = captions[0] self.caption_entry.delete(0, tk.END) self.caption_entry.insert(0, self.caption) #change the caption entry color to red self.caption_entry.configure(fg_color='red') def load_blip_model(self): self.blipSize = 384 blip_model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth' #check if options file exists if os.path.exists(os.path.join(self.captioner_folder, 'options.json')): with open(os.path.join(self.captioner_folder, 'options.json'), 'r') as f: self.nucleus_sampling = json.load(f)['nucleus_sampling'] self.q_factor = json.load(f)['q_factor'] self.min_length = json.load(f)['min_length'] else: self.nucleus_sampling = False self.q_factor = 1.0 self.min_length = 22 config_path = os.path.join(self.captioner_folder, "BLIP/configs/med_config.json") cache_folder = os.path.join(self.captioner_folder, "BLIP/cache") model_path = os.path.join(self.captioner_folder, "BLIP/models/model_base_caption_capfilt_large.pth") if not os.path.exists(cache_folder): os.makedirs(cache_folder) if not os.path.exists(model_path): print(f"Downloading BLIP to {cache_folder}") with requests.get(blip_model_url, stream=True) as session: session.raise_for_status() with open(model_path, 'wb') as f: for chunk in session.iter_content(chunk_size=1024): f.write(chunk) print('Download complete') else: print(f"Found BLIP model") import models.blip blip_decoder = models.blip.blip_decoder(pretrained=model_path, image_size=self.blipSize, vit='base', med_config=config_path) blip_decoder.eval() self.blip_decoder = blip_decoder.to(torch.device("cuda")) def batch_folder_mask(self): folder = '' try: # check if self.folder is set folder = self.folder except: pass dialog = BatchMaskWindow(self, folder) dialog.mainloop() def load_clip_seg_model(self): if self.clip_seg is None: self.clip_seg = ClipSeg() def open_folder(self,folder=None): if folder is None: self.folder = fd.askdirectory() else: self.folder = folder if self.folder == '': return self.output_folder = self.folder self.image_list = [os.path.join(self.folder, f) for f in os.listdir(self.folder) if _check_file_type(f) and not f.endswith('-masklabel.png') and not f.endswith('-depth.png')] #self.image_list.sort() #sort the image list alphabetically so that the images are in the same order every time self.image_list.sort(key=lambda x: x.lower()) self.image_count = len(self.image_list) if self.image_count == 0: tk.messagebox.showinfo("No Images", "No images found in the selected folder") return #update the image count label self.image_index = 0 self.image_count_label.configure(text=f'Image {self.image_index+1} of {self.image_count}') self.output_folder = self.folder self.load_image() self.caption_entry.focus_set() def draw_mask(self, event): if not self.enable_mask_editing.get(): return if event.widget != self.image_label.children["!label"]: return start_x = int(event.x / self.image_size[0] * self.PILimage.width) start_y = int(event.y / self.image_size[1] * self.PILimage.height) end_x = int(self.mask_draw_x / self.image_size[0] * self.PILimage.width) end_y = int(self.mask_draw_y / self.image_size[1] * self.PILimage.height) self.mask_draw_x = event.x self.mask_draw_y = event.y color = None if event.state & 0x0100 or event.num == 1: # left mouse button color = (255, 255, 255) elif event.state & 0x0400 or event.num == 3: # right mouse button color = (0, 0, 0) if color is not None: if self.PILmask is None: self.PILmask = Image.new('RGB', size=self.PILimage.size, color=(0, 0, 0)) draw = ImageDraw.Draw(self.PILmask) draw.line((start_x, start_y, end_x, end_y), fill=color, width=self.mask_draw_radius + self.mask_draw_radius + 1) draw.ellipse((start_x - self.mask_draw_radius, start_y - self.mask_draw_radius, start_x + self.mask_draw_radius, start_y + self.mask_draw_radius), fill=color, outline=None) draw.ellipse((end_x - self.mask_draw_radius, end_y - self.mask_draw_radius, end_x + self.mask_draw_radius, end_y + self.mask_draw_radius), fill=color, outline=None) self.compose_masked_image() self.display_image() def draw_mask_radius(self, event): if event.widget != self.image_label.children["!label"]: return delta = -np.sign(event.delta) * 5 self.mask_draw_radius += delta def compose_masked_image(self): np_image = np.array(self.PILimage).astype(np.float32) / 255.0 np_mask = np.array(self.PILmask).astype(np.float32) / 255.0 np_mask = np.clip(np_mask, 0.4, 1.0) np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) self.image = Image.fromarray(np_masked_image, mode='RGB') def display_image(self): #resize to fit 600x600 while maintaining aspect ratio width, height = self.image.size if width > height: new_width = 600 new_height = int(600 * height / width) else: new_height = 600 new_width = int(600 * width / height) self.image_size = (new_width, new_height) self.image = self.image.resize(self.image_size, Image.Resampling.LANCZOS) self.image = ctk.CTkImage(self.image, size=self.image_size) self.image_label.configure(image=self.image) def load_image(self): try: self.PILimage = Image.open(self.image_list[self.image_index]).convert('RGB') except: print(f'Error opening image {self.image_list[self.image_index]}') print('Logged path to bad_files.txt') #if bad_files.txt doesn't exist, create it if not os.path.exists('bad_files.txt'): with open('bad_files.txt', 'w') as f: f.write(self.image_list[self.image_index]+'\n') else: with open('bad_files.txt', 'a') as f: f.write(self.image_list[self.image_index]+'\n') return self.image = self.PILimage.copy() try: self.PILmask = None mask_filename = os.path.splitext(self.image_list[self.image_index])[0] + '-masklabel.png' if os.path.exists(mask_filename): self.PILmask = Image.open(mask_filename).convert('RGB') self.compose_masked_image() except Exception as e: print(f'Error opening mask for {self.image_list[self.image_index]}') print('Logged path to bad_files.txt') #if bad_files.txt doesn't exist, create it if not os.path.exists('bad_files.txt'): with open('bad_files.txt', 'w') as f: f.write(self.image_list[self.image_index]+'\n') else: with open('bad_files.txt', 'a') as f: f.write(self.image_list[self.image_index]+'\n') return self.display_image() self.caption_file_path = self.image_list[self.image_index] self.caption_file_name = os.path.basename(self.caption_file_path) self.caption_file_ext = os.path.splitext(self.caption_file_name)[1] self.caption_file_name_no_ext = os.path.splitext(self.caption_file_name)[0] self.caption_file = os.path.join(self.folder, self.caption_file_name_no_ext + '.txt') if os.path.isfile(self.caption_file) and self.auto_generate_caption.get() == False or os.path.isfile(self.caption_file) and self.auto_generate_caption.get() == True and self.auto_generate_caption_text_override.get() == True: with open(self.caption_file, 'r') as f: self.caption = f.read() self.caption_entry.delete(0, tk.END) self.caption_entry.insert(0, self.caption) self.caption_entry.configure(fg_color=ThemeManager.theme["CTkEntry"]["fg_color"]) self.use_blip = False elif os.path.isfile(self.caption_file) and self.auto_generate_caption.get() == True and self.auto_generate_caption_text_override.get() == False or os.path.isfile(self.caption_file)==False and self.auto_generate_caption.get() == True and self.auto_generate_caption_text_override.get() == True: self.use_blip = True self.caption_entry.delete(0, tk.END) elif os.path.isfile(self.caption_file) == False and self.auto_generate_caption.get() == False: self.caption_entry.delete(0, tk.END) return if self.use_blip and self.debug==False: tensor = transforms.Compose([ transforms.Resize((self.blipSize, self.blipSize), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) torch_image = tensor(self.PILimage).unsqueeze(0).to(torch.device("cuda")) if self.nucleus_sampling: captions = self.blip_decoder.generate(torch_image, sample=True, top_p=self.q_factor) else: captions = self.blip_decoder.generate(torch_image, sample=False, num_beams=16, min_length=self.min_length, \ max_length=48, repetition_penalty=self.q_factor) self.caption = captions[0] self.caption_entry.delete(0, tk.END) self.caption_entry.insert(0, self.caption) #change the caption entry color to red self.caption_entry.configure(fg_color='red') def save(self, event): self.save_caption() if self.enable_mask_editing.get(): self.save_mask() def save_mask(self): mask_filename = os.path.splitext(self.image_list[self.image_index])[0] + '-masklabel.png' if self.PILmask is not None: self.PILmask.save(mask_filename) def save_caption(self): self.caption = self.caption_entry.get() self.replace = self.replace_entry.get() self.replace_with = self.with_entry.get() self.suffix_var = self.suffix_entry.get() self.prefix = self.prefix_entry.get() #prepare the caption self.caption = self.caption.replace(self.replace, self.replace_with) if self.suffix_var.startswith(',') or self.suffix_var.startswith(' '): self.suffix_var = self.suffix_var else: self.suffix_var = ' ' + self.suffix_var if self.prefix != '': if self.prefix.endswith(' '): self.prefix = self.prefix[:-1] if not self.prefix.endswith(','): self.prefix = self.prefix+',' self.caption = self.prefix + ' ' + self.caption if self.caption.endswith(',') or self.caption.endswith('.'): self.caption = self.caption[:-1] self.caption = self.caption +', ' + self.suffix_var else: self.caption = self.caption + self.suffix_var self.caption = self.caption.strip() if self.output_folder != self.folder: outputFolder = self.output_folder else: outputFolder = self.folder if self.output_format == 'text': #text file with same name as image #image name #print('test') imgName = os.path.basename(self.image_list[self.image_index]) imgName = imgName[:imgName.rfind('.')] self.caption_file = os.path.join(outputFolder, imgName + '.txt') with open(self.caption_file, 'w') as f: f.write(self.caption) elif self.output_format == 'filename': #duplicate image with caption as file name #make sure self.caption doesn't contain any illegal characters illegal_chars = ['/', '\\', ':', '*', '?', '"', "'",'<', '>', '|', '.'] for char in illegal_chars: self.caption = self.caption.replace(char, '') self.PILimage.save(os.path.join(outputFolder, self.caption+'.png')) self.caption_entry.delete(0, tk.END) self.caption_entry.insert(0, self.caption) self.caption_entry.configure(fg_color='green') self.caption_entry.focus_force() def delete_word(self,event): ent = event.widget end_idx = ent.index(tk.INSERT) start_idx = ent.get().rfind(" ", None, end_idx) ent.selection_range(start_idx, end_idx) def prev_image(self, event): if self.image_index > 0: self.image_index -= 1 self.image_count_label.configure(text=f'Image {self.image_index+1} of {self.image_count}') self.load_image() self.caption_entry.focus_set() self.caption_entry.focus_force() def next_image(self, event): if self.image_index < len(self.image_list) - 1: self.image_index += 1 self.image_count_label.configure(text=f'Image {self.image_index+1} of {self.image_count}') self.load_image() self.caption_entry.focus_set() self.caption_entry.focus_force() def open_options(self): self.options_window = ctk.CTkToplevel(self) self.options_window.title("Options") self.options_window.geometry("320x550") #disable reszie self.options_window.resizable(False, False) self.options_window.focus_force() self.options_window.grab_set() self.options_window.transient(self) self.options_window.protocol("WM_DELETE_WINDOW", self.close_options) #add title label self.options_title_label = ctk.CTkLabel(self.options_window, text="Options",font=ctk.CTkFont(size=20, weight="bold")) self.options_title_label.pack(side="top", pady=5) #add an entry with a button to select a folder as output folder self.output_folder_label = ctk.CTkLabel(self.options_window, text="Output Folder") self.output_folder_label.pack(side="top", pady=5) self.output_folder_entry = ctk.CTkEntry(self.options_window) self.output_folder_entry.pack(side="top", fill="x", expand=False,padx=15, pady=5) self.output_folder_entry.insert(0, self.output_folder) self.output_folder_button = ctk.CTkButton(self.options_window, text="Select Folder", command=self.select_output_folder,fg_color=("gray75", "gray25")) self.output_folder_button.pack(side="top", pady=5) #add radio buttons to select the output format between text and filename self.output_format_label = ctk.CTkLabel(self.options_window, text="Output Format") self.output_format_label.pack(side="top", pady=5) self.output_format_var = tk.StringVar(self.options_window) self.output_format_var.set(self.output_format) self.output_format_text = ctk.CTkRadioButton(self.options_window, text="Text File", variable=self.output_format_var, value="text") self.output_format_text.pack(side="top", pady=5) self.output_format_filename = ctk.CTkRadioButton(self.options_window, text="File name", variable=self.output_format_var, value="filename") self.output_format_filename.pack(side="top", pady=5) #add BLIP settings section self.blip_settings_label = ctk.CTkLabel(self.options_window, text="BLIP Settings",font=ctk.CTkFont(size=20, weight="bold")) self.blip_settings_label.pack(side="top", pady=10) #add a checkbox to use nucleas sampling or not self.nucleus_sampling_var = tk.IntVar(self.options_window) self.nucleus_sampling_checkbox = ctk.CTkCheckBox(self.options_window, text="Use nucleus sampling", variable=self.nucleus_sampling_var) self.nucleus_sampling_checkbox.pack(side="top", pady=5) if self.debug: self.nucleus_sampling = 0 self.q_factor = 0.5 self.min_length = 10 self.nucleus_sampling_var.set(self.nucleus_sampling) #add a float entry to set the q factor self.q_factor_label = ctk.CTkLabel(self.options_window, text="Q Factor") self.q_factor_label.pack(side="top", pady=5) self.q_factor_entry = ctk.CTkEntry(self.options_window) self.q_factor_entry.insert(0, self.q_factor) self.q_factor_entry.pack(side="top", pady=5) #add a int entry to set the number minimum length self.min_length_label = ctk.CTkLabel(self.options_window, text="Minimum Length") self.min_length_label.pack(side="top", pady=5) self.min_length_entry = ctk.CTkEntry(self.options_window) self.min_length_entry.insert(0, self.min_length) self.min_length_entry.pack(side="top", pady=5) #add a horozontal radio button to select between None, ViT-L-14/openai, ViT-H-14/laion2b_s32b_b79k #self.model_label = ctk.CTkLabel(self.options_window, text="CLIP Interrogation") #self.model_label.pack(side="top") #self.model_var = tk.StringVar(self.options_window) #self.model_var.set(self.model) #self.model_none = tk.Radiobutton(self.options_window, text="None", variable=self.model_var, value="None") #self.model_none.pack(side="top") #self.model_vit_l_14 = tk.Radiobutton(self.options_window, text="ViT-L-14/openai", variable=self.model_var, value="ViT-L-14/openai") #self.model_vit_l_14.pack(side="top") #self.model_vit_h_14 = tk.Radiobutton(self.options_window, text="ViT-H-14/laion2b_s32b_b79k", variable=self.model_var, value="ViT-H-14/laion2b_s32b_b79k") #self.model_vit_h_14.pack(side="top") #add a save button self.save_button = ctk.CTkButton(self.options_window, text="Save", command=self.save_options, fg_color=("gray75", "gray25")) self.save_button.pack(side="top",fill='x',pady=10,padx=10) #all entries list entries = [self.output_folder_entry, self.q_factor_entry, self.min_length_entry] #bind the right click to all entries for entry in entries: entry.bind("", self.create_right_click_menu) self.options_file = os.path.join(self.captioner_folder, 'captioner_options.json') if os.path.isfile(self.options_file): with open(self.options_file, 'r') as f: self.options = json.load(f) self.output_folder_entry.delete(0, tk.END) self.output_folder_entry.insert(0, self.output_folder) self.output_format_var.set(self.options['output_format']) self.nucleus_sampling_var.set(self.options['nucleus_sampling']) self.q_factor_entry.delete(0, tk.END) self.q_factor_entry.insert(0, self.options['q_factor']) self.min_length_entry.delete(0, tk.END) self.min_length_entry.insert(0, self.options['min_length']) def load_options(self): self.options_file = os.path.join(self.captioner_folder, 'captioner_options.json') if os.path.isfile(self.options_file): with open(self.options_file, 'r') as f: self.options = json.load(f) #self.output_folder = self.folder #self.output_folder = self.options['output_folder'] if 'folder' in self.__dict__: self.output_folder = self.folder else: self.output_folder = '' self.output_format = self.options['output_format'] self.nucleus_sampling = self.options['nucleus_sampling'] self.q_factor = self.options['q_factor'] self.min_length = self.options['min_length'] else: #if self has folder, use it, otherwise use the current folder if 'folder' in self.__dict__ : self.output_folder = self.folder else: self.output_folder = '' self.output_format = "text" self.nucleus_sampling = False self.q_factor = 0.9 self.min_length =22 def save_options(self): self.output_folder = self.output_folder_entry.get() self.output_format = self.output_format_var.get() self.nucleus_sampling = self.nucleus_sampling_var.get() self.q_factor = float(self.q_factor_entry.get()) self.min_length = int(self.min_length_entry.get()) #save options to a file self.options_file = os.path.join(self.captioner_folder, 'captioner_options.json') with open(self.options_file, 'w') as f: json.dump({'output_folder': self.output_folder, 'output_format': self.output_format, 'nucleus_sampling': self.nucleus_sampling, 'q_factor': self.q_factor, 'min_length': self.min_length}, f) self.close_options() def select_output_folder(self): self.output_folder = fd.askdirectory() self.output_folder_entry.delete(0, tk.END) self.output_folder_entry.insert(0, self.output_folder) def close_options(self): self.options_window.destroy() self.caption_entry.focus_force() def create_right_click_menu(self, event): #create a menu self.menu = Menu(self, tearoff=0) #add commands to the menu self.menu.add_command(label="Cut", command=lambda: self.focus_get().event_generate("<>")) self.menu.add_command(label="Copy", command=lambda: self.focus_get().event_generate("<>")) self.menu.add_command(label="Paste", command=lambda: self.focus_get().event_generate("<>")) self.menu.add_command(label="Select All", command=lambda: self.focus_get().event_generate("<>")) #display the menu try: self.menu.tk_popup(event.x_root, event.y_root) finally: #make sure to release the grab (Tk 8.0a1 only) self.menu.grab_release() #progress bar class with cancel button class ProgressbarWithCancel(ctk.CTkToplevel): def __init__(self,max=None, **kw): super().__init__(**kw) self.title("Batching...") self.max = max self.possibleLabels = ['Searching for answers...',"I'm working, I promise.",'ARE THOSE TENTACLES?!','Weird data man...','Another one bites the dust' ,"I think it's a cat?" ,'Looking for the meaning of life', 'Dreaming of captions'] self.label = ctk.CTkLabel(self, text="Searching for answers...") self.label.pack(side="top", fill="x", expand=True,padx=10,pady=10) self.progress = ctk.CTkProgressBar(self, orientation="horizontal", mode="determinate") self.progress.pack(side="left", fill="x", expand=True,padx=10,pady=10) self.cancel_button = ctk.CTkButton(self, text="Cancel", command=self.cancel) self.cancel_button.pack(side="right",padx=10,pady=10) self.cancelled = False self.count_label = ctk.CTkLabel(self, text="0/{0}".format(self.max)) self.count_label.pack(side="right",padx=10,pady=10) def set_random_label(self): import random self.label["text"] = random.choice(self.possibleLabels) #pop from list #self.possibleLabels.remove(self.label["text"]) def cancel(self): self.cancelled = True def set_progress(self, value): self.progress.set(value) self.count_label.configure(text="{0}/{1}".format(int(value * self.max), self.max)) def get_progress(self): return self.progress.get def set_max(self, value): return value def get_max(self): return self.progress["maximum"] def is_cancelled(self): return self.cancelled #quit the progress bar window #run when imported as a module if __name__ == "__main__": #root = tk.Tk() app = ImageBrowser() app.mainloop()