import os import cv2 import numpy as np from PIL import Image, ImageOps, ImageDraw import platform import sys import tqdm import time from modules import shared, images from modules.shared import opts, cmd_opts if cmd_opts.deepdanbooru: import modules.deepbooru as deepbooru def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False, process_entropy_focus=False): try: if process_caption: shared.interrogator.load() if process_caption_deepbooru: db_opts = deepbooru.create_deepbooru_opts() db_opts[deepbooru.OPT_INCLUDE_RANKS] = False deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts) preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru, process_entropy_focus) finally: if process_caption: shared.interrogator.send_blip_to_ram() if process_caption_deepbooru: deepbooru.release_process() def preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False, process_entropy_focus=False): width = process_width height = process_height src = os.path.abspath(process_src) dst = os.path.abspath(process_dst) assert src != dst, 'same directory specified as source and destination' os.makedirs(dst, exist_ok=True) files = os.listdir(src) shared.state.textinfo = "Preprocessing..." shared.state.job_count = len(files) def save_pic_with_caption(image, index): caption = "" if process_caption: caption += shared.interrogator.generate_caption(image) if process_caption_deepbooru: if len(caption) > 0: caption += ", " caption += deepbooru.get_tags_from_process(image) filename_part = filename filename_part = os.path.splitext(filename_part)[0] filename_part = os.path.basename(filename_part) basename = f"{index:05}-{subindex[0]}-{filename_part}" image.save(os.path.join(dst, f"{basename}.png")) if len(caption) > 0: with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file: file.write(caption) subindex[0] += 1 def save_pic(image, index): save_pic_with_caption(image, index) if process_flip: save_pic_with_caption(ImageOps.mirror(image), index) for index, imagefile in enumerate(tqdm.tqdm(files)): subindex = [0] filename = os.path.join(src, imagefile) try: img = Image.open(filename).convert("RGB") except Exception: continue if shared.state.interrupted: break ratio = img.height / img.width is_tall = ratio > 1.35 is_wide = ratio < 1 / 1.35 processing_option_ran = False if process_split and is_tall: img = img.resize((width, height * img.height // img.width)) top = img.crop((0, 0, width, height)) save_pic(top, index) bot = img.crop((0, img.height - height, width, img.height)) save_pic(bot, index) processing_option_ran = True elif process_split and is_wide: img = img.resize((width * img.width // img.height, height)) left = img.crop((0, 0, width, height)) save_pic(left, index) right = img.crop((img.width - width, 0, img.width, height)) save_pic(right, index) processing_option_ran = True if process_entropy_focus and (is_tall or is_wide): if is_tall: img = img.resize((width, height * img.height // img.width)) else: img = img.resize((width * img.width // img.height, height)) x_focal_center, y_focal_center = image_central_focal_point(img, width, height) # take the focal point and turn it into crop coordinates that try to center over the focal # point but then get adjusted back into the frame y_half = int(height / 2) x_half = int(width / 2) x1 = x_focal_center - x_half if x1 < 0: x1 = 0 elif x1 + width > img.width: x1 = img.width - width y1 = y_focal_center - y_half if y1 < 0: y1 = 0 elif y1 + height > img.height: y1 = img.height - height x2 = x1 + width y2 = y1 + height crop = [x1, y1, x2, y2] focal = img.crop(tuple(crop)) save_pic(focal, index) processing_option_ran = True if not processing_option_ran: img = images.resize_image(1, img, width, height) save_pic(img, index) shared.state.nextjob() def image_central_focal_point(im, target_width, target_height): focal_points = [] focal_points.extend( image_focal_points(im) ) fp_entropy = image_entropy_point(im, target_width, target_height) fp_entropy['weight'] = len(focal_points) + 1 # about half of the weight to entropy focal_points.append(fp_entropy) weight = 0.0 x = 0.0 y = 0.0 for focal_point in focal_points: weight += focal_point['weight'] x += focal_point['x'] * focal_point['weight'] y += focal_point['y'] * focal_point['weight'] avg_x = round(x // weight) avg_y = round(y // weight) return avg_x, avg_y def image_focal_points(im): grayscale = im.convert("L") # naive attempt at preventing focal points from collecting at watermarks near the bottom gd = ImageDraw.Draw(grayscale) gd.rectangle([0, im.height*.9, im.width, im.height], fill="#999") np_im = np.array(grayscale) points = cv2.goodFeaturesToTrack( np_im, maxCorners=50, qualityLevel=0.04, minDistance=min(grayscale.width, grayscale.height)*0.05, useHarrisDetector=False, ) if points is None: return [] focal_points = [] for point in points: x, y = point.ravel() focal_points.append({ 'x': x, 'y': y, 'weight': 1.0 }) return focal_points def image_entropy_point(im, crop_width, crop_height): img = im.copy() # just make it easier to slide the test crop with images oriented the same way if (img.size[0] < img.size[1]): portrait = True img = img.rotate(90, expand=1) e_max = 0 crop_current = [0, 0, crop_width, crop_height] crop_best = crop_current while crop_current[2] < img.size[0]: crop = img.crop(tuple(crop_current)) e = image_entropy(crop) if (e_max < e): e_max = e crop_best = list(crop_current) crop_current[0] += 4 crop_current[2] += 4 x_mid = int((crop_best[2] - crop_best[0])/2) y_mid = int((crop_best[3] - crop_best[1])/2) return { 'x': x_mid, 'y': y_mid, 'weight': 1.0 } def image_entropy(im): # greyscale image entropy band = np.asarray(im.convert("L")) hist, _ = np.histogram(band, bins=range(0, 256)) hist = hist[hist > 0] return -np.log2(hist / hist.sum()).sum()