stable-diffusion-webui

1
from PIL import Image
2

3
from modules import scripts_postprocessing, ui_components
4
import gradio as gr
5

6

7
def center_crop(image: Image, w: int, h: int):
8
    iw, ih = image.size
9
    if ih / h < iw / w:
10
        sw = w * ih / h
11
        box = (iw - sw) / 2, 0, iw - (iw - sw) / 2, ih
12
    else:
13
        sh = h * iw / w
14
        box = 0, (ih - sh) / 2, iw, ih - (ih - sh) / 2
15
    return image.resize((w, h), Image.Resampling.LANCZOS, box)
16

17

18
def multicrop_pic(image: Image, mindim, maxdim, minarea, maxarea, objective, threshold):
19
    iw, ih = image.size
20
    err = lambda w, h: 1 - (lambda x: x if x < 1 else 1 / x)(iw / ih / (w / h))
21
    wh = max(((w, h) for w in range(mindim, maxdim + 1, 64) for h in range(mindim, maxdim + 1, 64)
22
              if minarea <= w * h <= maxarea and err(w, h) <= threshold),
23
             key=lambda wh: (wh[0] * wh[1], -err(*wh))[::1 if objective == 'Maximize area' else -1],
24
             default=None
25
             )
26
    return wh and center_crop(image, *wh)
27

28

29
class ScriptPostprocessingAutosizedCrop(scripts_postprocessing.ScriptPostprocessing):
30
    name = "Auto-sized crop"
31
    order = 4020
32

33
    def ui(self):
34
        with ui_components.InputAccordion(False, label="Auto-sized crop") as enable:
35
            gr.Markdown('Each image is center-cropped with an automatically chosen width and height.')
36
            with gr.Row():
37
                mindim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension lower bound", value=384, elem_id="postprocess_multicrop_mindim")
38
                maxdim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension upper bound", value=768, elem_id="postprocess_multicrop_maxdim")
39
            with gr.Row():
40
                minarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area lower bound", value=64 * 64, elem_id="postprocess_multicrop_minarea")
41
                maxarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area upper bound", value=640 * 640, elem_id="postprocess_multicrop_maxarea")
42
            with gr.Row():
43
                objective = gr.Radio(["Maximize area", "Minimize error"], value="Maximize area", label="Resizing objective", elem_id="postprocess_multicrop_objective")
44
                threshold = gr.Slider(minimum=0, maximum=1, step=0.01, label="Error threshold", value=0.1, elem_id="postprocess_multicrop_threshold")
45

46
        return {
47
            "enable": enable,
48
            "mindim": mindim,
49
            "maxdim": maxdim,
50
            "minarea": minarea,
51
            "maxarea": maxarea,
52
            "objective": objective,
53
            "threshold": threshold,
54
        }
55

56
    def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, mindim, maxdim, minarea, maxarea, objective, threshold):
57
        if not enable:
58
            return
59

60
        cropped = multicrop_pic(pp.image, mindim, maxdim, minarea, maxarea, objective, threshold)
61
        if cropped is not None:
62
            pp.image = cropped
63
        else:
64
            print(f"skipped {pp.image.width}x{pp.image.height} image (can't find suitable size within error threshold)")
65