fix image orientation and tune preprocessing

cli/modules/preview.py

@@ -1,6 +1,10 @@
 #!/bin/env python
+"""
+create preview images from embeddings
+"""
 import os
 import io
+import sys
 import json
 import base64
 from pathlib import Path
@@ -11,12 +15,10 @@ from util import Map, log
 from sdapi import getsync, postsync
 from grid import grid
 
-# masks = ['preview-face.jpg', 'preview-body.jpg']
-mask = 'preview-body.jpg'
-template = 'photo of "{name}", {suffix}, high detailed, skin texture, facing camera, 135mm, shot on dslr, 4k, modelshoot style'
+template = 'photo of "{name}", {suffix}, high detailed, skin texture, lookin forward, facing camera, 135mm, shot on dslr, 4k, modelshoot style'
 img2img_options = Map({
     'prompt': None,
-    'negative_prompt': '',
+    'negative_prompt': 'cartoon, drawing, cgi, sketch, comic, disfigured, deformed',
     'init_images': [],
     'sampler_name': 'DPM2 Karras',
     'batch_size': 4,
@@ -40,12 +42,14 @@ def create_preview(name: str, suffix: str):
     options = getsync('/sdapi/v1/options')
     cmdflags = getsync('/sdapi/v1/cmd-flags')    
     img2img_options['prompt'] = template.format(name = name, suffix = suffix)
-    log.info({ 'preview prompt': img2img_options['prompt'] })
     log.debug({ 'preview options': img2img_options })
-    mask_path = os.path.join(os.path.dirname(getsourcefile(lambda:0)), mask)
     if len(img2img_options['init_images']) == 0:
         for i in range(img2img_options.batch_size):
-            img2img_options['init_images'].append(encode(mask_path))
+            mask = os.path.join(os.path.dirname(getsourcefile(lambda:0)), 'preview'+ str(i+1) +'.jpg')
+            if (not os.path.isfile(mask)):
+                log.error({ 'preview': 'missing preview mask' })
+                return
+            img2img_options['init_images'].append(encode(mask))
     data = postsync('/sdapi/v1/img2img', img2img_options)
     if 'error' in data:
         log.error({ 'preview': data['error'], 'reason': data['reason'] })
@@ -65,5 +69,13 @@ def create_preview(name: str, suffix: str):
 if __name__ == "__main__":
     log.info({ 'preview': 'start' })
     cmdflags = getsync('/sdapi/v1/cmd-flags')
-    for f in Path(cmdflags.embeddings_dir).glob('*.pt'):
-        create_preview(f.stem, 'person')
+    sys.argv.pop(0)
+    if len(sys.argv) == 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)
+    log.info({ 'preview embeddings': len(files) })
+    for f in files:
+        name = Path(f).stem
+        create_preview(name, 'person')

cli/modules/process.py

@@ -43,9 +43,9 @@ params = Map({
     'square_images': True, # should output images be squared
     'blur_samplesize': 60, # sample size to use for blur detection
     'face_score': 0.7, # min face detection score
-    'face_pad': 0.05, # pad face image percentage
+    '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.2, # max score for face blur detection
+    'face_blur_score': 1.4, # max score for face blur 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
@@ -117,7 +117,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.warning({ 'extract face': 'out of frame' })
+        log.info({ '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
@@ -130,7 +130,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.warning({ 'extract face': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
+        log.info({ 'extract face': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
         return None, True
     cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
 
@@ -144,14 +144,14 @@ def extract_face(img):
 
     blur = detect_blur(squared)
     if blur > params.face_blur_score:
-        log.warning({ 'extract face': 'blur check fail', 'blur': blur })
+        log.info({ 'extract face': 'blur check fail', 'blur': blur })
         return None, True
     else:
-        log.info({ 'extract face blur': blur })
+        log.debug({ 'extract face blur': blur })
 
     similarity = detect_simmilar(squared)
     if similarity > params.similarity_score:
-        log.warning({ 'extract face': 'similarity check fail', 'score': similarity })
+        log.info({ 'extract face': 'similarity check fail', 'score': round(similarity, 2) })
         return None, True
 
     return squared, True
@@ -172,7 +172,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.warning({ 'extract body': 'insufficient body parts', 'detected': len(x) })
+        log.info({ 'extract body': '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
@@ -183,7 +183,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.warning({ 'extract body': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
+        log.info({ 'extract body': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
         return None, True
     cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
 
@@ -197,14 +197,14 @@ def extract_body(img):
 
     blur = detect_blur(squared)
     if blur > params.body_blur_score:
-        log.warning({ 'extract body': 'blur check fail', 'blur': blur })
+        log.info({ 'extract body': 'blur check fail', 'blur': blur })
         return None, True
     else:
-        log.info({ 'extract body blur': blur })
+        log.debug({ 'extract body blur': blur })
 
     similarity = detect_simmilar(squared)
     if similarity > params.similarity_score:
-        log.warning({ 'extract body': 'similarity check fail', 'score': similarity })
+        log.info({ 'extract body': 'similarity check fail', 'score': similarity })
         return None, True
 
     return squared, True
@@ -252,17 +252,19 @@ def process_file(f: str, dst: str = None):
         log.error({ 'image': f, 'error': err })
         return
 
+    image = ImageOps.exif_transpose(image) # rotate image according to EXIF orientation
+
     if image.width < 512 or image.height < 512:
-        log.warning({ 'skip low resolution': [image.width, image.height], 'file': f })
+        log.info({ 'skip low resolution': [image.width, image.height], 'file': f })
         return
-    log.info({ 'resolution': [image.width, image.height], 'mp': round((image.width * image.height) / 1024 / 1024, 1) })
+    log.debug({ 'resolution': [image.width, image.height], 'mp': round((image.width * image.height) / 1024 / 1024, 1) })
 
     face, ok = extract_face(image)
     if face is not None:
         fn = save(face, f, 'face')
         log.info({ 'extract face': fn })
     else:
-        log.warning({ 'no face': f })
+        log.debug({ 'no face': f })
 
     if not ok:
         return
@@ -272,7 +274,7 @@ def process_file(f: str, dst: str = None):
         fn = save(body, f, 'body')
         log.info({ 'extract body': fn })   
     else:
-        log.warning({ 'no body': f })
+        log.debug({ 'no body': f })
 
 
 def process_images(src: str, dst: str, args = None):
@@ -285,7 +287,7 @@ def process_images(src: str, dst: str, args = None):
         log.error({ 'process': 'not a folder', 'src': src })
     else:
         if os.path.isdir(dst) and params.clear_dst:
-            log.warning({ 'clear dst': dst })
+            log.info({ 'clear dst': 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)

cli/train.py

@@ -78,7 +78,8 @@ async def preprocess_cleanup(params):
     for f in Path(params.dst).glob('*.txt'):
         f.unlink()
     try:
-        Path(params.dst).rmdir()
+        if os.path.isdir(params.dst):
+            Path(params.dst).rmdir()
     except Exception as err:
         log.warning({ 'preprocess cleanup': params.dst, 'error': err })
 

requirements.txt

@@ -2,6 +2,7 @@ accelerate
 basicsr
 clean-fid
 einops
+fastapi
 filetype
 font-roboto
 fonts