update

cli/modules/ffmpeg.py

@@ -16,13 +16,15 @@ def probe(src: str):
     cmd = f"ffprobe -hide_banner -loglevel 0 -print_format json -show_format -show_streams {src}"
     result = subprocess.run(cmd, shell = True, capture_output = True, text = True, check = True)
     data = json.loads(result.stdout)
-    i = [x for x in data['streams'] if x["codec_type"] == "video"][0]
+    stream = [x for x in data['streams'] if x["codec_type"] == "video"][0]
+    format = data['format'] if 'format' in data else {}
+    res = {**stream, **format}
     video = Map({
-        'codec': i['codec_name']+'/'+i['codec_tag_string'],
-        'resolution': [int(i['width']), int(i['height'])],
-        'duration': float(i['duration']),
-        'frames': int(i['nb_frames']),
-        'bitrate': round(float(i['bit_rate']) / 1024),
+        'codec': res.get('codec_name', 'unknown') + '/' + res.get('codec_tag_string', ''),
+        'resolution': [int(res.get('width', 0)), int(res.get('height', 0))],
+        'duration': float(res.get('duration', 0)),
+        'frames': int(res.get('nb_frames', 0)),
+        'bitrate': round(float(res.get('bit_rate', 0)) / 1024),
     })
     return video
 

cli/modules/grid.py

@@ -77,6 +77,7 @@ if __name__ == '__main__':
     parser.add_argument("--width", type = int, default = 0, required = False, help = "fixed grid width")
     parser.add_argument("--height", type = int, default = 0, required = False, help = "fixed grid height")
     parser.add_argument("--border", type = int, default = 0, required = False, help = "image border")
+    parser.add_argument('--nolabels', default = False, action='store_true', help = "do not print image labels")
     parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information")
     parser.add_argument('output', type = str)
     parser.add_argument('input', type = str, nargs = '*')
@@ -106,7 +107,8 @@ if __name__ == '__main__':
             # img.verify()
             images.append(img)
             fp = Path(file)
-            labels.append(fp.stem)
+            if not params.nolabels:
+                labels.append(fp.stem)
     # log.info({ 'folder': path.parent, 'labels': labels })
     if len(images) > 0:
         image = grid(

cli/modules/process.py

@@ -1,17 +1,33 @@
 #!/bin/env python
 """
-process images
+process people images
+- check image resolution
+- runs detection of face and body
+- extracts crop and performs checks:
+    - visible: is face or body detected
+    - 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
+    - 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
+- if image passes checks  
+  image padded and saved as extracted image
+- body requires that face is detected and in-frame,  
+  but does not have to pass all other checks as body performs its own checks
+- runs clip interrogation on extracted images to generate filewords
 """
 import os
 import sys
 import io
-import shutil
+import filetype
 import base64
 import pathlib
 
 import numpy as np
 import mediapipe as mp
 from PIL import Image, ImageOps
+from skimage.metrics import structural_similarity as ssim
 
 from util import log, Map
 from sdapi import postsync
@@ -26,7 +42,7 @@ params = Map({
     'target_size': 512, # target resolution
     'square_images': True, # should output images be squared
     'blur_samplesize': 60, # sample size to use for blur detection
-    'face_score': 0.6, # min face detection score
+    'face_score': 0.7, # min face detection score
     'face_pad': 0.05, # pad face image percentage
     'face_model': 1, # which face model to use 0/close-up 1/standard
     'face_blur_score': 1.2, # max score for face blur detection
@@ -40,6 +56,8 @@ params = Map({
     'segmentation_body': False, # segmentation enabled
     'segmentation_model': 0, # segmentation model 0/general 1/landscape
     'segmentation_background': (192, 192, 192), # segmentation background color
+    'similarity_score': 0.6, # maximum similarity score before image is discarded
+    'similarity_size': 64, # base similarity detection on reduced images
     'interrogate_model': 'clip' # interrogate model
 })
 
@@ -58,6 +76,20 @@ def detect_blur(image):
     return mean
 
 
+images = []
+def detect_simmilar(image):
+    img = image.resize((params.similarity_size, params.similarity_size))
+    img = ImageOps.grayscale(img)
+    data = np.array(img)
+    similarity = 0
+    for i in images:
+        val = ssim(data, i, data_range=255, channel_axis=None, gradient=False, full=False)
+        if val > similarity:
+            similarity = val
+    images.append(data)
+    return similarity
+
+
 def segmentation(image):
     with mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=params.segmentation_model) as selfie_segmentation:
         data = np.array(image)
@@ -78,11 +110,15 @@ def extract_face(img):
     scale = max(img.size[0], img.size[1]) / params.target_size
     resized = img.copy()
     resized.thumbnail((params.target_size, params.target_size), Image.HAMMING)
+
     with mp.solutions.face_detection.FaceDetection(min_detection_confidence=params.face_score, model_selection=params.face_model) as face:
         results = face.process(np.array(resized))
     if results.detections is None:
         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.warning({ 'extract face': 'out of frame' })
+        return None, False
     x = (box.xmin - params.face_pad / 2) * resized.width
     y = (box.ymin - params.face_pad / 2)* resized.height
     w = (box.width + params.face_pad) * resized.width
@@ -97,6 +133,7 @@ def extract_face(img):
         log.warning({ 'extract face': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
         return None, True
     cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
+
     if params.square_images:
         squared = Image.new('RGB', (params.target_size, params.target_size))
         squared.paste(cropped, (0, 0))
@@ -104,6 +141,7 @@ def extract_face(img):
            squared = segmentation(squared)
     else:
         squared = cropped
+
     blur = detect_blur(squared)
     if blur > params.face_blur_score:
         log.warning({ 'extract face': 'blur check fail', 'blur': blur })
@@ -111,6 +149,11 @@ def extract_face(img):
     else:
         log.info({ 'extract face blur': blur })
 
+    similarity = detect_simmilar(squared)
+    if similarity > params.similarity_score:
+        log.warning({ 'extract face': 'similarity check fail', 'score': similarity })
+        return None, True
+
     return squared, True
 
 
@@ -143,6 +186,7 @@ def extract_body(img):
         log.warning({ 'extract body': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
         return None, True
     cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
+
     if params.square_images:
         squared = Image.new('RGB', (params.target_size, params.target_size))
         squared.paste(cropped, (0, 0))
@@ -150,12 +194,19 @@ def extract_body(img):
            squared = segmentation(squared)
     else:
         squared = cropped
+
     blur = detect_blur(squared)
     if blur > params.body_blur_score:
         log.warning({ 'extract body': 'blur check fail', 'blur': blur })
         return None, True
     else:
         log.info({ 'extract body blur': blur })
+
+    similarity = detect_simmilar(squared)
+    if similarity > params.similarity_score:
+        log.warning({ 'extract body': 'similarity check fail', 'score': similarity })
+        return None, True
+
     return squared, True
 
 
@@ -235,7 +286,9 @@ def process_images(src: str, dst: str, args = None):
     else:
         if os.path.isdir(dst) and params.clear_dst:
             log.warning({ 'clear dst': dst })
-            shutil.rmtree(dst) 
+            i = [os.path.join(dst, f) for f in os.listdir(dst) if os.path.isfile(os.path.join(dst, f)) and filetype.is_image(os.path.join(dst, f))]
+            for f in i:
+                os.remove(f)
         pathlib.Path(dst).mkdir(parents=True, exist_ok=True)
         for root, _sub_dirs, files in os.walk(src):
             for f in files:
@@ -248,6 +301,7 @@ if __name__ == '__main__':
     dst = sys.argv.pop(0)
     params.dst = dst
     log.info({ 'processing': params })
+    pathlib.Path(dst).mkdir(parents=True, exist_ok=True)
     for loc in sys.argv:
         if os.path.isfile(loc):
             process_file(loc, dst)

cli/train.json

@@ -2,9 +2,9 @@
   "training_model": "sd-v15-runwayml.ckpt",
   "extract_video": {
     "rate": 0,
-    "fps": 2,
-    "skipstart": 0,
-    "skipend": 0
+    "fps": 5,
+    "vstart": 0,
+    "vend": 0
   },
   "create_embedding": {
       "name": "test",

cli/train.md

@@ -22,8 +22,7 @@
 
 ## Prompt
 
-    a medium shot photo of "dreamkelly", extremely detailed 8k wallpaper, intricate, high detail, dramatic, modelshoot style
-
+    a medium shot photo of "kelly", extremely detailed 8k wallpaper, intricate, high detail, dramatic, modelshoot style
 
 ## Gen
 

cli/train.py

@@ -265,9 +265,21 @@ async def train(params):
     log.debug({ 'train start' })
     args.train_embedding.embedding_name = params.name
     args.train_embedding.data_root = args.preprocess.process_dst
+    imgs = [f for f in os.listdir(args.preprocess.process_dst) if os.path.isfile(os.path.join(args.preprocess.process_dst, f)) and filetype.is_image(os.path.join(args.preprocess.process_dst, f))]
+    if len(imgs) == 0:
+        log.error({ 'train no input images in folder': args.preprocess.process_dst })
+        return
+    if params.grad == -1:
+        grad = (len(imgs) // args.train_embedding.batch_size)
+        args.train_embedding.gradient_step = max(grad, 30)
+        log.info({ 'dynamic gradient step': args.train_embedding.gradient_step })
+        if params.steps == -1:
+            args.train_embedding.steps = 5000 // args.train_embedding.gradient_step
+            log.info({ 'dynamic steps': args.train_embedding.steps })
     log.info({ 'train embedding': {
         'name': params.name,
         'source': args.preprocess.process_dst,
+        'images': len(imgs),
         'steps': args.train_embedding.steps,
         'batch': args.train_embedding.batch_size,
         'gradient-step': args.train_embedding.gradient_step,
@@ -376,14 +388,14 @@ async def main():
     parser.add_argument("--src", type = str, required = True, help = "source image folder or movie file")
     parser.add_argument("--init", type = str, default = "person", required = False, help = "initialization class, default: %(default)s")
     parser.add_argument("--dst", type = str, default = "/tmp", required = False, help = "destination image folder for processed images, default: %(default)s")
-    parser.add_argument("--steps", type = int, default = 250, required = False, help = "training steps, default: %(default)s")
+    parser.add_argument("--steps", type = int, default = -1, required = False, help = "training steps, default: %(default)s")
     parser.add_argument("--vectors", type = int, default = -1, required = False, help = "number of vectors per token, default: dynamic based on number of input images")
     parser.add_argument("--batch", type = int, default = 1, required = False, help = "batch size, default: %(default)s")
     parser.add_argument("--rate", type = str, default = "", required = False, help = "learning rate, default: dynamic")
     parser.add_argument("--rstart", type = float, default = 0.01, required = False, help = "starting learn rate if using dynamic rate, default: %(default)s")
     parser.add_argument("--rend", type = float, default = 0.0001, required = False, help = "ending learn rate if using dynamic rate, default: %(default)s")
     parser.add_argument("--rdescend", type = float, default = 2, required = False, help = "learn rate descend power when using dynamic rate, default: %(default)s")
-    parser.add_argument("--grad", type = int, default = 20, required = False, help = "accumulate gradient over n images, default: : %(default)s")
+    parser.add_argument("--grad", type = int, default = -1, required = False, help = "accumulate gradient over n images, default: : %(default)s")
     parser.add_argument("--type", type = str, default = 'subject', required = False, help = "training type: subject/style/unknown, default: %(default)s")
     parser.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing embedding, default: %(default)s")
     parser.add_argument("--vstart", type = float, default = 0, required = False, help = "if processing video skip first n seconds, default: %(default)s")