diff --git a/TODO.md b/TODO.md index c8b80344a..bbc109645 100644 --- a/TODO.md +++ b/TODO.md @@ -88,13 +88,19 @@ Cool stuff that is not integrated anywhere... - must run `./automatic.sh install` - note: this is quite a big one so some testing is reccomended after upgrade - non-trivial ui updates -- added **brightness dynamic range** check to `process.py` -- new script: `watermark.py` +- renamed scripts in `cli/modules` to be more descriptive + if you're using old script names, update them + for example, `ffmpeg.py` is now `video-extract.py` +- updated script `process.py` + - new **brightness dynamic range** check + - new **preview** mode to run all checks but without saving images plus print a summary at the end +- new script: `image-watermark.py` - optionally strip exif from images - add invisible watermark to images which persists even if user modifies image so we can always track it -- new script: `palette.py` +- new script: `palette-extract.py` - creates color palette wheel from image - not finished +- updated `embedding-preview.py` so it can skip existing previews or overwrite them - expose variation seed in main ui - integrated seed travel functionality into core - integrated `pix2pix` functionality to standard `img2img` workflow diff --git a/cli/generate.py b/cli/generate.py index 498864cfd..7c267ed2a 100755 --- a/cli/generate.py +++ b/cli/generate.py @@ -25,6 +25,7 @@ import pathlib import secrets import time import sys +import importlib from random import randrange from PIL import Image @@ -330,8 +331,8 @@ async def main(): dynamic = prompt(params) if params.beautify: try: - from modules.promptist import beautify # pylint: disable=import-outside-toplevel - sd.generate.prompt = beautify(dynamic) + promptist = importlib.import_module('modules.promptist') + sd.generate.prompt = promptist.beautify(dynamic) except Exception as e: log.error({ 'beautify': e }) scheduler = sampler(params, options) diff --git a/cli/bench.py b/cli/modules/bench.py similarity index 100% rename from cli/bench.py rename to cli/modules/bench.py diff --git a/cli/modules/preview.py b/cli/modules/embedding-preview.py similarity index 72% rename from cli/modules/preview.py rename to cli/modules/embedding-preview.py index b634bd9ff..3d5763e14 100755 --- a/cli/modules/preview.py +++ b/cli/modules/embedding-preview.py @@ -7,6 +7,7 @@ import io import sys import json import base64 +import argparse from pathlib import Path from PIL import Image from inspect import getsourcefile @@ -45,7 +46,7 @@ def create_preview(name: str, suffix: str): log.debug({ 'preview options': img2img_options }) if len(img2img_options['init_images']) == 0: for i in range(img2img_options.batch_size): - mask = os.path.join(os.path.dirname(getsourcefile(lambda:0)), 'preview'+ str(i+1) +'.jpg') + mask = os.path.join(os.path.dirname(getsourcefile(lambda:0)), 'preview-template'+ str(i+1) +'.jpg') if (not os.path.isfile(mask)): log.error({ 'preview': 'missing preview mask' }) return @@ -69,12 +70,32 @@ def create_preview(name: str, suffix: str): if __name__ == "__main__": log.info({ 'preview': 'start' }) cmdflags = getsync('/sdapi/v1/cmd-flags') - sys.argv.pop(0) - if len(sys.argv) == 0: + + parser = argparse.ArgumentParser(description = 'generate embeddings previews') + parser.add_argument('--overwrite', default = False, action='store_true', help = 'overwrite existing previews') + parser.add_argument('input', type=str, nargs='*') + params = parser.parse_args() + + if len(params.input) == 0: files = list(Path(cmdflags.embeddings_dir).glob('*.pt')) else: - files = list(os.path.join(cmdflags.embeddings_dir, a + '.pt') for a in sys.argv if os.path.isfile(os.path.join(cmdflags.embeddings_dir, a + '.pt'))) - files.sort(key=os.path.getctime, reverse=True) + files = list(os.path.join(cmdflags.embeddings_dir, a + '.pt') for a in params.input if os.path.isfile(os.path.join(cmdflags.embeddings_dir, a + '.pt'))) + candidates = [str(f) for f in files] + candidates.sort(key=os.path.getctime, reverse=True) + + files = [] + for f in candidates: + fn = f.replace('.pt', '.preview.png') + if os.path.isfile(f.replace('.pt', '.preview.png')): + if params.overwrite: + log.info({ 'preview add': fn }) + files.append(f) + else: + log.info({ 'preview skip': fn }) + else: + log.info({ 'preview add': fn }) + files.append(f) + log.info({ 'preview embeddings': len(files) }) for f in files: name = Path(f).stem diff --git a/cli/modules/watermark.py b/cli/modules/image-watermark.py similarity index 100% rename from cli/modules/watermark.py rename to cli/modules/image-watermark.py diff --git a/cli/modules/lossrate.jpg b/cli/modules/lossrate.jpg deleted file mode 100644 index 8e310b74e..000000000 Binary files a/cli/modules/lossrate.jpg and /dev/null differ diff --git a/cli/models.py b/cli/modules/models-preview.py similarity index 100% rename from cli/models.py rename to cli/modules/models-preview.py diff --git a/cli/modules/palette.py b/cli/modules/palette-extract.py similarity index 100% rename from cli/modules/palette.py rename to cli/modules/palette-extract.py diff --git a/cli/modules/preview1.jpg b/cli/modules/preview-template1.jpg similarity index 100% rename from cli/modules/preview1.jpg rename to cli/modules/preview-template1.jpg diff --git a/cli/modules/preview2.jpg b/cli/modules/preview-template2.jpg similarity index 100% rename from cli/modules/preview2.jpg rename to cli/modules/preview-template2.jpg diff --git a/cli/modules/preview3.jpg b/cli/modules/preview-template3.jpg similarity index 100% rename from cli/modules/preview3.jpg rename to cli/modules/preview-template3.jpg diff --git a/cli/modules/preview4.jpg b/cli/modules/preview-template4.jpg similarity index 100% rename from cli/modules/preview4.jpg rename to cli/modules/preview-template4.jpg diff --git a/cli/modules/process.py b/cli/modules/process.py index 05581cbc8..89e352b99 100755 --- a/cli/modules/process.py +++ b/cli/modules/process.py @@ -8,6 +8,7 @@ process people images - in frame: for face based on box, for body based on number of visible keypoints - resolution: is cropped image still of sufficient resolution - blur: is image sharp enough + - dynamic range: is image bright enough - similarity: compares image to all previously processed images to see if its unique enough - images are resized and optionally squared - face additionally runs through semantic segmentation to remove background @@ -18,11 +19,12 @@ process people images - runs clip interrogation on extracted images to generate filewords """ import os -import sys import io import math import base64 import pathlib +import argparse +import logging import filetype import numpy as np @@ -48,14 +50,14 @@ params = Map({ 'face_pad': 0.07, # pad face image percentage 'face_model': 1, # which face model to use 0/close-up 1/standard 'face_blur_score': 1.5, # max score for face blur detection - 'face_range_score': 0.5, # min score for face dynamic range detection + 'face_range_score': 0.3, # min score for face dynamic range detection 'body_score': 0.9, # min body detection score 'body_visibility': 0.5, # min visibility score for each detected body part 'body_parts': 15, # min number of detected body parts with sufficient visibility 'body_pad': 0.2, # pad body image percentage 'body_model': 2, # body model to use 0/low 1/medium 2/high 'body_blur_score': 1.8, # max score for body blur detection - 'body_range_score': 0.5, # min score for body dynamic range detection + 'body_range_score': 0.3, # min score for body dynamic range detection 'segmentation_face': True, # segmentation enabled 'segmentation_body': False, # segmentation enabled 'segmentation_model': 0, # segmentation model 0/general 1/landscape @@ -140,7 +142,7 @@ def extract_face(img): return None, False box = results.detections[0].location_data.relative_bounding_box if box.xmin < 0 or box.ymin < 0 or (box.width - box.xmin) > 1 or (box.height - box.ymin) > 1: - log.info({ 'extract face': 'out of frame' }) + log.info({ 'process face skip': 'out of frame' }) return None, False x = (box.xmin - params.face_pad / 2) * resized.width y = (box.ymin - params.face_pad / 2)* resized.height @@ -153,7 +155,7 @@ def extract_face(img): square = [max(square[0], 0), max(square[1], 0), min(square[2], img.width), min(square[3], img.height)] cropped = img.crop(tuple(square)) if cropped.size[0] < params.target_size and cropped.size[1] < params.target_size: - log.info({ 'extract face': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] }) + log.info({ 'process face skip': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] }) return None, True cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING) @@ -167,21 +169,21 @@ def extract_face(img): blur = detect_blur(squared) if blur > params.face_blur_score: - log.info({ 'extract face': 'blur check fail', 'blur': blur }) + log.info({ 'process face skip': 'blur check fail', 'blur': blur }) return None, True else: - log.debug({ 'extract face blur': blur }) + log.debug({ 'process face blur': blur }) range = detect_dynamicrange(squared) if range < params.face_range_score: - log.info({ 'extract face': 'dynamic range check fail', 'range': range }) + log.info({ 'process face skip': 'dynamic range check fail', 'range': range }) return None, True else: - log.debug({ 'extract face dynamic range': range }) + log.debug({ 'process face dynamic range': range }) similarity = detect_simmilar(squared) if similarity > params.similarity_score: - log.info({ 'extract face': 'similarity check fail', 'score': round(similarity, 2) }) + log.info({ 'process face skip': 'similarity check fail', 'score': round(similarity, 2) }) return None, True return squared, True @@ -202,7 +204,7 @@ def extract_body(img): x = [resized.width * (i.x - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility] y = [resized.height * (i.y - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility] if len(x) < params.body_parts: - log.info({ 'extract body': 'insufficient body parts', 'detected': len(x) }) + log.info({ 'process body skip': 'insufficient body parts', 'detected': len(x) }) return None, True w = max(x) - min(x) + resized.width * params.body_pad h = max(y) - min(y) + resized.height * params.body_pad @@ -213,7 +215,7 @@ def extract_body(img): square = [max(square[0], 0), max(square[1], 0), min(square[2], img.width), min(square[3], img.height)] cropped = img.crop(tuple(square)) if cropped.size[0] < params.target_size and cropped.size[1] < params.target_size: - log.info({ 'extract body': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] }) + log.info({ 'process body skip': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] }) return None, True cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING) @@ -227,21 +229,21 @@ def extract_body(img): blur = detect_blur(squared) if blur > params.body_blur_score: - log.info({ 'extract body': 'blur check fail', 'blur': blur }) + log.info({ 'process body skip': 'blur check fail', 'blur': blur }) return None, True else: - log.debug({ 'extract body blur': blur }) + log.debug({ 'process body blur': blur }) range = detect_dynamicrange(squared) if range < params.body_range_score: - log.info({ 'extract body': 'dynamic range check fail', 'range': range }) + log.info({ 'process body skip': 'dynamic range check fail', 'range': range }) return None, True else: - log.debug({ 'extract body dynamic range': range }) + log.debug({ 'process body dynamic range': range }) similarity = detect_simmilar(squared) if similarity > params.similarity_score: - log.info({ 'extract body': 'similarity check fail', 'score': similarity }) + log.info({ 'process body skip': 'similarity check fail', 'score': similarity }) return None, True return squared, True @@ -268,7 +270,7 @@ def interrogate(img, fn): i = {} -def process_file(f: str, dst: str = None): +def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = False): def save(img, f, what): i[what] = i.get(what, 0) + 1 if dst is None: @@ -278,8 +280,10 @@ def process_file(f: str, dst: str = None): base = os.path.basename(f).split('.')[0] fn = os.path.join(dir, str(i[what]).rjust(3, '0') + '-' + what + '-' + base + '.jpg') # log.debug({ 'save': fn }) - img.save(fn) - interrogate(img, fn) + if not preview: + img.save(fn) + if not offline: + interrogate(img, fn) return fn log.info({ 'processing': f }) @@ -287,12 +291,12 @@ def process_file(f: str, dst: str = None): image = Image.open(f) except Exception as err: log.error({ 'image': f, 'error': err }) - return + return 0, 0 image = ImageOps.exif_transpose(image) # rotate image according to EXIF orientation if image.width < 512 or image.height < 512: - log.info({ 'skip low resolution': [image.width, image.height], 'file': f }) + log.info({ 'process skip': 'low resolution', 'resolution': [image.width, image.height] }) return log.debug({ 'resolution': [image.width, image.height], 'mp': round((image.width * image.height) / 1024 / 1024, 1) }) @@ -337,15 +341,28 @@ def process_images(src: str, dst: str, args = None): if __name__ == '__main__': # log.setLevel(logging.DEBUG) - sys.argv.pop(0) - dst = sys.argv.pop(0) - params.dst = dst + parser = argparse.ArgumentParser(description = 'image watermarking') + parser.add_argument('--output', type=str, required=True, help='folder to store images') + parser.add_argument('--preview', default=False, action='store_true', help = "run processing but do not store results") + parser.add_argument('--offline', default=False, action='store_true', help = "run only processing steps that do not require running server") + parser.add_argument('--debug', default=False, action='store_true', help = "enable debug logging") + parser.add_argument('input', type=str, nargs='*') + args = parser.parse_args() + params.dst = args.output + if args.debug: + log.setLevel(logging.DEBUG) + log.debug({ 'debug': True }) log.info({ 'processing': params }) - pathlib.Path(dst).mkdir(parents=True, exist_ok=True) - for loc in sys.argv: + if not os.path.exists(params.dst) and not args.preview: + pathlib.Path(params.dst).mkdir(parents=True, exist_ok=True) + files = [] + for loc in args.input: if os.path.isfile(loc): - process_file(loc, dst) + files.append(loc) elif os.path.isdir(loc): - for root, _sub_dirs, files in os.walk(loc): - for f in files: - process_file(os.path.join(root, f), dst) + for root, _sub_dirs, dir in os.walk(loc): + for f in dir: + files.append(os.path.join(root, f)) + for f in files: + process_file(f, params.dst, args.preview, args.offline) + log.info({ 'processed': i, 'inputs': len(files) }) diff --git a/cli/modules/ideas.py b/cli/modules/prompt-ideas.py similarity index 100% rename from cli/modules/ideas.py rename to cli/modules/prompt-ideas.py diff --git a/cli/modules/promptist.py b/cli/modules/prompt-promptist.py similarity index 100% rename from cli/modules/promptist.py rename to cli/modules/prompt-promptist.py diff --git a/cli/modules/losschart.py b/cli/modules/train-losschart.py similarity index 100% rename from cli/modules/losschart.py rename to cli/modules/train-losschart.py diff --git a/cli/modules/lossrate.py b/cli/modules/train-lossrate.py similarity index 100% rename from cli/modules/lossrate.py rename to cli/modules/train-lossrate.py diff --git a/cli/modules/ffmpeg.py b/cli/modules/video-extract.py similarity index 92% rename from cli/modules/ffmpeg.py rename to cli/modules/video-extract.py index 89e86c431..5f9bc9f7a 100755 --- a/cli/modules/ffmpeg.py +++ b/cli/modules/video-extract.py @@ -62,7 +62,6 @@ def extract(src: str, dst: str, rate: float = 0.015, fps: float = 0, start = 0, if __name__ == "__main__": parser = argparse.ArgumentParser(description="ffmpeg pipeline") - parser.add_argument("command", choices = ["extract", "animate"]) parser.add_argument("--input", type = str, required = True, help="input") parser.add_argument("--output", type = str, required = True, help="output") parser.add_argument("--rate", type = float, default = 0, required = False, help="extraction change rate threshold") @@ -70,5 +69,4 @@ if __name__ == "__main__": parser.add_argument("--skipstart", type = float, default = 1, required = False, help="skip time from start of video") parser.add_argument("--skipend", type = float, default = 1, required = False, help="skip time to end of video") params = parser.parse_args() - if params.command == "extract": - extract(src = params.input, dst = params.output, rate = params.rate, fps = params.fps, start = params.skipstart, end = params.skipend) + extract(src = params.input, dst = params.output, rate = params.rate, fps = params.fps, start = params.skipstart, end = params.skipend) diff --git a/cli/modules/dynamotest.py b/cli/random/dynamotest.py similarity index 100% rename from cli/modules/dynamotest.py rename to cli/random/dynamotest.py diff --git a/cli/train.py b/cli/train.py index 1f5c48cf5..7bfb1ec02 100755 --- a/cli/train.py +++ b/cli/train.py @@ -18,6 +18,7 @@ import os import sys import time import json +import importlib from pathlib import Path, PurePath import filetype @@ -25,14 +26,13 @@ from PIL import Image sys.path.append(os.path.join(os.path.dirname(__file__), 'modules')) from modules.util import Map, log, set_logfile -from modules.losschart import plot -from modules.ffmpeg import extract -from modules.lossrate import gen_loss_rate_str from modules.sdapi import close, get, interrupt, post, progress, session from modules.process import process_images from modules.grid import grid -from modules.preview import create_preview - +create_preview = importlib.import_module('modules.embedding-preview').create_preview +plot = importlib.import_module('modules.train-losschart').plot +extract = importlib.import_module('modules.video-extract').extract +gen_loss_rate_str = importlib.import_module('modules.train-lossrate').gen_loss_rate_str images = [] args = {}