diff --git a/cli/modules/ffmpeg.py b/cli/modules/ffmpeg.py index 650bfaaeb..d33d97e4c 100755 --- a/cli/modules/ffmpeg.py +++ b/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 diff --git a/cli/modules/grid.py b/cli/modules/grid.py index 36c245f82..83798e619 100755 --- a/cli/modules/grid.py +++ b/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( diff --git a/cli/modules/process.py b/cli/modules/process.py index 00f92f072..cdfd6c58d 100755 --- a/cli/modules/process.py +++ b/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) diff --git a/cli/train.json b/cli/train.json index 94f4cab7d..d6b3050bc 100644 --- a/cli/train.json +++ b/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", diff --git a/cli/train.md b/cli/train.md index fc5bad090..076fe2098 100644 --- a/cli/train.md +++ b/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 diff --git a/cli/train.py b/cli/train.py index a941f7482..adc309920 100755 --- a/cli/train.py +++ b/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")