restructure scripts

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

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)

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

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) })

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)

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 = {}