mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 07:31:01 +02:00
Vladimir Mandic
238 lines
8.6 KiB
Python
Executable File
238 lines
8.6 KiB
Python
Executable File
#!/bin/env python
|
|
"""
|
|
process images
|
|
"""
|
|
import os
|
|
import sys
|
|
import io
|
|
import shutil
|
|
import base64
|
|
import pathlib
|
|
|
|
import numpy as np
|
|
import mediapipe as mp
|
|
from PIL import Image, ImageOps
|
|
|
|
from util import log, Map
|
|
from sdapi import postsync
|
|
|
|
|
|
params = Map({
|
|
'src': '', # source folder
|
|
'dst': '', # destination folder
|
|
'clear_dst': True, # remove all files from destination at the start
|
|
'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_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
|
|
'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.6, # max score for body blur detection
|
|
'interrogate_model': 'clip' # interrogate model
|
|
})
|
|
|
|
|
|
def detect_blur(image):
|
|
# based on <https://github.com/karthik9319/Blur-Detection/>
|
|
bw = ImageOps.grayscale(image)
|
|
cx, cy = image.size[0] // 2, image.size[1] // 2
|
|
fft = np.fft.fft2(bw)
|
|
fftShift = np.fft.fftshift(fft)
|
|
fftShift[cy - params.blur_samplesize: cy + params.blur_samplesize, cx - params.blur_samplesize: cx + params.blur_samplesize] = 0
|
|
fftShift = np.fft.ifftshift(fftShift)
|
|
recon = np.fft.ifft2(fftShift)
|
|
magnitude = np.log(np.abs(recon))
|
|
mean = round(np.mean(magnitude), 2)
|
|
return mean
|
|
|
|
|
|
def extract_face(img):
|
|
if img.mode == 'RGBA':
|
|
img = img.convert('RGB')
|
|
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
|
|
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
|
|
h = (box.height + params.face_pad) * resized.height
|
|
cx = x + w / 2
|
|
cy = y + h / 2
|
|
l = max(w, h) / 2
|
|
square = [scale * (cx - l), scale * (cy - l), scale * (cx + l), scale * (cy + l)]
|
|
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]] })
|
|
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))
|
|
else:
|
|
squared = cropped
|
|
blur = detect_blur(squared)
|
|
if blur > params.face_blur_score:
|
|
log.warning({ 'extract face': 'blur check fail', 'blur': blur })
|
|
return None, True
|
|
else:
|
|
log.info({ 'extract face blur': blur })
|
|
|
|
return squared, True
|
|
|
|
|
|
def extract_body(img):
|
|
if img.mode == 'RGBA':
|
|
img = img.convert('RGB')
|
|
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.pose.Pose(static_image_mode=True, min_detection_confidence=params.body_score, model_complexity=params.body_model) as pose:
|
|
results = pose.process(np.array(resized))
|
|
if results.pose_landmarks is None:
|
|
return None, False
|
|
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) })
|
|
return None, True
|
|
w = max(x) - min(x) + resized.width * params.body_pad
|
|
h = max(y) - min(y) + resized.height * params.body_pad
|
|
cx = min(x) + w / 2
|
|
cy = min(y) + h / 2
|
|
l = max(w, h) / 2
|
|
square = [scale * (cx - l), scale * (cy - l), scale * (cx + l), scale * (cy + l)]
|
|
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]] })
|
|
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))
|
|
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 })
|
|
return squared, True
|
|
|
|
|
|
def interrogate(img, fn):
|
|
def encode(f):
|
|
with io.BytesIO() as stream:
|
|
img.save(stream, 'JPEG')
|
|
values = stream.getvalue()
|
|
encoded = base64.b64encode(values).decode()
|
|
return encoded
|
|
|
|
if params.interrogate_model is None or params.interrogate_model == '':
|
|
return
|
|
json = Map({ 'image': encode(img), 'model': params.interrogate_model })
|
|
res = postsync('/sdapi/v1/interrogate', json)
|
|
caption = res.caption if 'caption' in res else ''
|
|
log.info({ 'interrogate': caption })
|
|
file = fn.replace('.jpg', '.txt')
|
|
f = open(file, 'w')
|
|
f.write(caption)
|
|
f.close()
|
|
|
|
|
|
i = {}
|
|
def process_file(f: str, dst: str = None):
|
|
def save(img, f, what):
|
|
i[what] = i.get(what, 0) + 1
|
|
if dst is None:
|
|
dir = os.path.dirname(f)
|
|
else:
|
|
dir = dst
|
|
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)
|
|
return fn
|
|
|
|
log.info({ 'processing': f })
|
|
try:
|
|
image = Image.open(f)
|
|
except Exception as err:
|
|
log.error({ 'image': f, 'error': err })
|
|
return
|
|
|
|
if image.width < 512 or image.height < 512:
|
|
log.warning({ '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) })
|
|
|
|
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 })
|
|
|
|
if not ok:
|
|
return
|
|
|
|
body, ok = extract_body(image)
|
|
if body is not None:
|
|
fn = save(body, f, 'body')
|
|
log.info({ 'extract body': fn })
|
|
else:
|
|
log.warning({ 'no body': f })
|
|
|
|
|
|
def process_images(src: str, dst: str, args = None):
|
|
params.src = src
|
|
params.dst = dst
|
|
if args is not None:
|
|
params.update(args)
|
|
log.info({ 'processing': params })
|
|
if not os.path.isdir(src):
|
|
log.error({ 'process': 'not a folder', 'src': src })
|
|
else:
|
|
if os.path.isdir(dst) and params.clear_dst:
|
|
log.warning({ 'clear dst': dst })
|
|
shutil.rmtree(dst)
|
|
pathlib.Path(dst).mkdir(parents=True, exist_ok=True)
|
|
for root, _sub_dirs, files in os.walk(src):
|
|
for f in files:
|
|
process_file(os.path.join(root, f), dst)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# log.setLevel(logging.DEBUG)
|
|
sys.argv.pop(0)
|
|
dst = sys.argv.pop(0)
|
|
params.dst = dst
|
|
log.info({ 'processing': params })
|
|
for loc in sys.argv:
|
|
if os.path.isfile(loc):
|
|
process_file(loc, dst)
|
|
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)
|
|
|
|
"""
|
|
- add interrogate on save
|
|
- write final stats
|
|
- create final grid
|
|
"""
|