restructure scripts
@@ -0,0 +1,128 @@
|
||||
#!/bin/env python
|
||||
"""
|
||||
sd api txt2img benchmark
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
from PIL import Image
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), 'modules'))
|
||||
import modules.sdapi as sdapi
|
||||
from modules.util import Map, log
|
||||
|
||||
|
||||
options = Map({
|
||||
'restore_faces': False,
|
||||
'prompt': 'photo of two dice on a table',
|
||||
'negative_prompt': 'foggy, blurry',
|
||||
'steps': 20,
|
||||
'batch_size': 1,
|
||||
'n_iter': 1,
|
||||
'seed': -1,
|
||||
'sampler_name': 'Euler a',
|
||||
'cfg_scale': 0,
|
||||
'width': 512,
|
||||
'height': 512
|
||||
})
|
||||
|
||||
|
||||
# batch = [1, 1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128]
|
||||
batch = [1, 1, 2, 4, 8, 12, 16]
|
||||
oom = 0
|
||||
|
||||
|
||||
async def txt2img():
|
||||
t0 = time.perf_counter()
|
||||
data = {}
|
||||
try:
|
||||
data = await sdapi.post('/sdapi/v1/txt2img', options)
|
||||
except:
|
||||
return -1
|
||||
if 'error' in data:
|
||||
return -1
|
||||
if 'info' in data:
|
||||
info = Map(json.loads(data['info']))
|
||||
else:
|
||||
return 0
|
||||
log.debug({ 'info': info })
|
||||
for i in range(len(data['images'])):
|
||||
data['images'][i] = Image.open(io.BytesIO(base64.b64decode(data['images'][i].split(',',1)[0])))
|
||||
log.debug({ 'image': data['images'][i].size })
|
||||
t1 = time.perf_counter()
|
||||
return t1 - t0
|
||||
|
||||
|
||||
def memstats():
|
||||
mem = sdapi.getsync('/sdapi/v1/memory')
|
||||
cpu = mem.get('ram', 'unavailable')
|
||||
gpu = mem.get('cuda', 'unavailable')
|
||||
if 'active' in gpu:
|
||||
gpu['session'] = gpu.pop('active')
|
||||
if 'reserved' in gpu:
|
||||
gpu.pop('allocated')
|
||||
gpu.pop('reserved')
|
||||
gpu.pop('inactive')
|
||||
if 'events' in gpu:
|
||||
global oom # pylint: disable=global-statement
|
||||
oom = gpu['events']['oom']
|
||||
gpu.pop('events')
|
||||
return cpu, gpu
|
||||
|
||||
|
||||
def gb(val: float):
|
||||
return round(val / 1024 / 1024 / 1024, 2)
|
||||
|
||||
|
||||
async def main():
|
||||
log.info({ 'benchmark': { 'batch-sizes': batch } })
|
||||
sdapi.quiet = True
|
||||
await sdapi.session()
|
||||
await sdapi.interrupt()
|
||||
opts = await sdapi.get('/sdapi/v1/options')
|
||||
opts = Map(opts)
|
||||
log.info({ 'options': {
|
||||
'resolution': [options.width, options.height],
|
||||
'model': opts.sd_model_checkpoint,
|
||||
'vae': opts.sd_vae,
|
||||
'hypernetwork': opts.sd_hypernetwork,
|
||||
'sampler': options.sampler_name,
|
||||
'clip-stop': opts.CLIP_stop_at_last_layers,
|
||||
'preview': opts.show_progress_every_n_steps
|
||||
} })
|
||||
cpu, gpu = memstats()
|
||||
log.info({ 'system': { 'cpu': cpu, 'gpu': gpu }})
|
||||
for i in range(len(batch)):
|
||||
if oom > 0:
|
||||
continue
|
||||
options['batch_size'] = batch[i]
|
||||
ts = await txt2img()
|
||||
if ts > 0:
|
||||
await asyncio.sleep(0)
|
||||
cpu, gpu = memstats()
|
||||
if i == 0:
|
||||
log.info({ 'warmup': round(ts, 2) })
|
||||
else:
|
||||
peak = gpu['session']['peak'] if 'session' in gpu else 0
|
||||
log.info({ 'batch': batch[i], 'its': round(options.steps / (ts / batch[i]), 2), 'img': round(ts / batch[i], 2), 'wall': round(ts, 2), 'peak': gb(peak), 'oom': oom > 0 })
|
||||
else:
|
||||
await asyncio.sleep(10)
|
||||
cpu, gpu = memstats()
|
||||
log.info({ 'batch': batch[i], 'result': 'error', 'gpu': gpu, 'oom': oom > 0 })
|
||||
break
|
||||
if oom > 0:
|
||||
log.info({ 'benchmark': 'ended with oom so you should probably restart your automatic server now' })
|
||||
await sdapi.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
log.warning({ 'interrupted': 'keyboard request' })
|
||||
sdapi.interruptsync()
|
||||
@@ -1,103 +0,0 @@
|
||||
#!/bin/env python
|
||||
"""
|
||||
Test Torch Dynamo functionality and backends
|
||||
"""
|
||||
import json
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision.models import resnet18
|
||||
|
||||
|
||||
print('torch:', torch.__version__)
|
||||
try:
|
||||
import torch._dynamo as dynamo # must be imported explicitly or namespace is not found
|
||||
except Exception as err:
|
||||
print('torch without dynamo support', err)
|
||||
|
||||
|
||||
N_ITERS = 20
|
||||
torch._dynamo.config.verbose=True
|
||||
warnings.filterwarnings('ignore', category=UserWarning) # disable those for now as many backends reports tons
|
||||
# torch.set_float32_matmul_precision('high') # enable to test in fp32
|
||||
|
||||
|
||||
def timed(fn): # returns the result of running `fn()` and the time it took for `fn()` to run in ms using CUDA events
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
result = fn()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return result, start.elapsed_time(end)
|
||||
|
||||
|
||||
def generate_data(b):
|
||||
return (
|
||||
torch.randn(b, 3, 128, 128).to(torch.float32).cuda(),
|
||||
torch.randint(1000, (b,)).cuda(),
|
||||
)
|
||||
|
||||
|
||||
def init_model():
|
||||
return resnet18().to(torch.float32).cuda()
|
||||
|
||||
|
||||
def eval(mod, inp):
|
||||
return mod(inp)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# first pass, dynamo is going to be slower as it compiles
|
||||
model = init_model()
|
||||
inp = generate_data(16)[0]
|
||||
|
||||
# repeat test
|
||||
results = {}
|
||||
times = []
|
||||
print('eager initial eval:', timed(lambda: eval(model, inp))[1])
|
||||
for i in range(N_ITERS):
|
||||
inp = generate_data(16)[0]
|
||||
_res, time = timed(lambda: eval(model, inp))
|
||||
times.append(time)
|
||||
results['default'] = np.median(times)
|
||||
|
||||
print('dynamo available backends:', dynamo.list_backends())
|
||||
for backend in dynamo.list_backends():
|
||||
try:
|
||||
torch._dynamo.reset() # required before changing backends
|
||||
eval_dyn = dynamo.optimize(backend)(eval)
|
||||
print('dynamo initial eval:', backend, timed(lambda: eval_dyn(model, inp))[1])
|
||||
times = []
|
||||
for i in range(N_ITERS):
|
||||
inp = generate_data(16)[0]
|
||||
_res, time = timed(lambda: eval_dyn(model, inp))
|
||||
times.append(time)
|
||||
results[backend] = np.median(times)
|
||||
except Exception as err:
|
||||
lines = str(err).split('\n')
|
||||
print('dyanmo backend failed:', backend, lines[0]) # print just first error line as backtraces can be quite long
|
||||
results[backend] = 'error'
|
||||
|
||||
# print stats
|
||||
print(json.dumps(results, indent = 4))
|
||||
|
||||
|
||||
"""
|
||||
Reference: <https://github.com/pytorch/pytorch/blob/4f4b62e4a255708e928445b6502139d5962974fa/docs/source/dynamo/get-started.rst>
|
||||
Training & Inference backends:
|
||||
dynamo.optimize("inductor") - Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton kernels
|
||||
dynamo.optimize("aot_nvfuser") - nvFuser with AotAutograd
|
||||
dynamo.optimize("aot_cudagraphs") - cudagraphs with AotAutograd
|
||||
Inference-only backends:
|
||||
dynamo.optimize("ofi") - Uses Torchscript optimize_for_inference
|
||||
dynamo.optimize("fx2trt") - Uses Nvidia TensorRT for inference optimizations
|
||||
dynamo.optimize("onnxrt") - Uses ONNXRT for inference on CPU/GPU
|
||||
Results:
|
||||
"default": 4.247040033340454,
|
||||
"ofi": 3.820032000541687,
|
||||
"aot_cudagraphs": 6.460927963256836,
|
||||
"inductor": RuntimeError: CUDA: Error- no device
|
||||
"fx2trt": ImportError: libtorch_cuda_cu.so: cannot open shared object file: No such file or directory
|
||||
"""
|
||||
@@ -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
|
||||
|
Before Width: | Height: | Size: 163 KiB |
@@ -0,0 +1,129 @@
|
||||
#!/bin/env python
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
import argparse
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), 'modules'))
|
||||
from generate import sd, generate
|
||||
from modules.util import Map, log
|
||||
from modules.sdapi import get, post, close
|
||||
from modules.grid import grid
|
||||
|
||||
|
||||
default = 'sd-v15-runwayml.ckpt [cc6cb27103]'
|
||||
embeddings = ['blonde', 'bruntette', 'sexy', 'naked', 'mia', 'lin', 'kelly', 'hanna', 'rreid-random-v0']
|
||||
exclude = ['sd-v20', 'sd-v21', 'inpainting', 'pix2pix']
|
||||
prompt = "photo of beautiful woman <embedding>, photograph, posing, pose, high detailed, intricate, elegant, sharp focus, skin texture, looking forward, facing camera, 135mm, shot on dslr, canon 5d, 4k, modelshoot style, cinematic lighting"
|
||||
options = Map({
|
||||
'generate': {
|
||||
'restore_faces': True,
|
||||
'prompt': '',
|
||||
'negative_prompt': 'digital art, cgi, render, foggy, blurry, blurred, duplicate, ugly, mutilated, mutation, mutated, out of frame, bad anatomy, disfigured, deformed, censored, low res, low resolution, watermark, text, poorly drawn face, poorly drawn hands, signature',
|
||||
'steps': 30,
|
||||
'batch_size': 4,
|
||||
'n_iter': 1,
|
||||
'seed': -1,
|
||||
'sampler_name': 'DPM2 Karras',
|
||||
'cfg_scale': 7,
|
||||
'width': 512,
|
||||
'height': 512
|
||||
},
|
||||
'paths': {
|
||||
"root": "/mnt/c/Users/mandi/OneDrive/Generative/Generate",
|
||||
"generate": "image",
|
||||
"upscale": "upscale",
|
||||
"grid": "grid"
|
||||
},
|
||||
'options': {
|
||||
"sd_model_checkpoint": "sd-v15-runwayml",
|
||||
"sd_vae": "vae-ft-mse-840000-ema-pruned.ckpt"
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
async def models(params):
|
||||
global sd
|
||||
data = await get('/sdapi/v1/sd-models')
|
||||
all = [m['title'] for m in data]
|
||||
models = []
|
||||
excluded = []
|
||||
for m in all: # loop through all registered models
|
||||
ok = True
|
||||
for e in exclude: # check if model is excluded
|
||||
if e in m:
|
||||
excluded.append(m)
|
||||
ok = False
|
||||
break
|
||||
if ok:
|
||||
short = m.split(' [')[0]
|
||||
short = short.replace('.ckpt', '').replace('.safetensors', '')
|
||||
models.append(short)
|
||||
if len(params.input) > 0: # check if model is included in cmd line
|
||||
filtered = []
|
||||
for m in params.input:
|
||||
if m in models:
|
||||
filtered.append(m)
|
||||
else:
|
||||
log.error({ 'model not found': m })
|
||||
return
|
||||
models = filtered
|
||||
log.info({ 'models preview' })
|
||||
log.info({ 'models': len(models), 'excluded': len(excluded) })
|
||||
log.info({ 'embeddings': embeddings })
|
||||
cmdflags = await get('/sdapi/v1/cmd-flags')
|
||||
opt = await get('/sdapi/v1/options')
|
||||
if params.output != '':
|
||||
dir = params.output
|
||||
else:
|
||||
dir = os.path.abspath(os.path.join(cmdflags['hypernetwork_dir'], '..', 'Stable-diffusion'))
|
||||
log.info({ 'output directory': dir })
|
||||
log.info({ 'total jobs': len(models) * len(embeddings) * options.generate.batch_size, 'per-model': len(embeddings) * options.generate.batch_size })
|
||||
log.info(json.dumps(options, indent=2))
|
||||
for model in models:
|
||||
fn = os.path.join(dir, model + '.png')
|
||||
if os.path.exists(fn) and len(params.input) == 0: # if model preview exists and not manually included
|
||||
log.info({ 'model preview exists': model })
|
||||
continue
|
||||
log.info({ 'model load': model })
|
||||
|
||||
opt['sd_model_checkpoint'] = model
|
||||
await post('/sdapi/v1/options', opt)
|
||||
opt = await get('/sdapi/v1/options')
|
||||
images = []
|
||||
labels = []
|
||||
t0 = time.time()
|
||||
for embedding in embeddings:
|
||||
options.generate.prompt = prompt.replace('<embedding>', f'\"{embedding}\"')
|
||||
log.info({ 'model generating': model, 'embedding': embedding, 'prompt': options.generate.prompt })
|
||||
data = await generate(options = options, quiet=True)
|
||||
if 'image' in data:
|
||||
for img in data['image']:
|
||||
images.append(img)
|
||||
labels.append(embedding)
|
||||
else:
|
||||
log.error({ 'model': model, 'embedding': embedding, 'error': data })
|
||||
t1 = time.time()
|
||||
image = grid(images = images, labels = labels, border = 8)
|
||||
image.save(fn)
|
||||
t = t1 - t0
|
||||
its = 1.0 * options.generate.steps * len(images) / t
|
||||
log.info({ 'model preview created': model, 'image': fn, 'images': len(images), 'grid': [image.width, image.height], 'time': round(t, 2), 'its': round(its, 2) })
|
||||
|
||||
opt = await get('/sdapi/v1/options')
|
||||
if opt['sd_model_checkpoint'] != default:
|
||||
log.info({ 'model set default': default })
|
||||
opt['sd_model_checkpoint'] = default
|
||||
await post('/sdapi/v1/options', opt)
|
||||
|
||||
await close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description = 'generate model previews')
|
||||
parser.add_argument('--output', type = str, default = '', required = False, help = 'output directory')
|
||||
parser.add_argument('input', type = str, nargs = '*')
|
||||
params = parser.parse_args()
|
||||
asyncio.run(models(params))
|
||||
|
Before Width: | Height: | Size: 7.1 KiB After Width: | Height: | Size: 7.1 KiB |
|
Before Width: | Height: | Size: 8.0 KiB After Width: | Height: | Size: 8.0 KiB |
|
Before Width: | Height: | Size: 7.6 KiB After Width: | Height: | Size: 7.6 KiB |
|
Before Width: | Height: | Size: 9.1 KiB After Width: | Height: | Size: 9.1 KiB |
@@ -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) })
|
||||
|
||||
@@ -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)
|
||||