mirror of
https://github.com/vladmandic/automatic
synced 2026-08-29 16:41:01 +02:00
Vladimir Mandic
102 lines
4.2 KiB
Python
102 lines
4.2 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import gradio as gr
|
|
from PIL import Image
|
|
import modules.scripts as scripts
|
|
from modules import images, processing, shared
|
|
from modules.processing import Processed
|
|
from modules.shared import opts, state
|
|
|
|
|
|
class Script(scripts.Script):
|
|
def title(self):
|
|
return "HDR"
|
|
|
|
def show(self, is_img2img):
|
|
return True
|
|
|
|
def ui(self, is_img2img):
|
|
with gr.Row():
|
|
gr.HTML("<span>  High Dynamic Range</span><br>")
|
|
with gr.Row():
|
|
save_hdr = gr.Checkbox(label="Save HDR image", value=True)
|
|
hdr_range = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.65, label='HDR range')
|
|
with gr.Row():
|
|
is_tonemap = gr.Checkbox(label="Enable tonemap", value=False)
|
|
gamma = gr.Slider(minimum=0, maximum=2, step=0.05, value=1.0, label='Gamma', visible=False)
|
|
with gr.Row():
|
|
scale = gr.Slider(minimum=0, maximum=2, step=0.05, value=1.0, label='Scale', visible=False)
|
|
saturation = gr.Slider(minimum=0, maximum=2, step=0.05, value=1.0, label='Saturation', visible=False)
|
|
is_tonemap.change(fn=self.change_tonemap, inputs=[is_tonemap], outputs=[gamma, scale, saturation])
|
|
return [hdr_range, save_hdr, is_tonemap, gamma, scale, saturation]
|
|
|
|
def change_tonemap(self, is_tonemap):
|
|
return [gr.update(visible=is_tonemap), gr.update(visible=is_tonemap), gr.update(visible=is_tonemap)]
|
|
|
|
def merge(self, imgs: list, is_tonemap: bool, gamma, scale, saturation):
|
|
shared.log.info(f'HDR: merge images={len(imgs)} tonemap={is_tonemap} sgamma={gamma} scale={scale} saturation={saturation}')
|
|
imgs_np = [np.asarray(img).astype(np.uint8) for img in imgs]
|
|
|
|
align = cv2.createAlignMTB()
|
|
align.process(imgs_np, imgs_np)
|
|
|
|
# cv2.createMergeRobertson()
|
|
# cv2.createMergeDebevec()
|
|
merge = cv2.createMergeMertens()
|
|
hdr = merge.process(imgs_np)
|
|
|
|
# cv2.createTonemapDrago()
|
|
# cv2.createTonemapReinhard()
|
|
if is_tonemap:
|
|
tonemap = cv2.createTonemapMantiuk(gamma, scale, saturation)
|
|
hdr = tonemap.process(hdr)
|
|
|
|
ldr = np.clip(hdr * 255, 0, 255).astype(np.uint8)
|
|
hdr = np.clip(hdr * 65535, 0, 65535).astype(np.uint16)
|
|
hdr = cv2.cvtColor(hdr, cv2.COLOR_BGR2RGB)
|
|
return hdr, ldr
|
|
|
|
def run(self, p, hdr_range, save_hdr, is_tonemap, gamma, scale, saturation): # pylint: disable=arguments-differ
|
|
if shared.sd_model_type != 'sd' and shared.sd_model_type != 'sdxl':
|
|
shared.log.error(f'HDR: incorrect base model: {shared.sd_model.__class__.__name__}')
|
|
return
|
|
p.extra_generation_params = {
|
|
"HDR range": hdr_range,
|
|
}
|
|
shared.log.info(f'HDR: range={hdr_range}')
|
|
processing.fix_seed(p)
|
|
imgs = []
|
|
info = ''
|
|
for i in range(3):
|
|
p.n_iter = 1
|
|
p.batch_size = 1
|
|
p.do_not_save_grid = True
|
|
p.hdr_brightness = (i - 1) * (2.0 * hdr_range)
|
|
p.hdr_mode = 0
|
|
p.task_args['seed'] = p.seed
|
|
processed: processing.Processed = processing.process_images(p)
|
|
imgs += processed.images
|
|
if i == 1:
|
|
info = processed.info
|
|
if state.interrupted:
|
|
break
|
|
|
|
if len(imgs) > 1:
|
|
hdr, ldr = self.merge(imgs, is_tonemap, gamma, scale, saturation)
|
|
img = Image.fromarray(ldr)
|
|
imgs.insert(0, img)
|
|
if save_hdr:
|
|
fn, _txt = images.save_image(img, shared.opts.outdir_save, "", p.seed, p.prompt, opts.grid_format, info=processed.info, p=p)
|
|
fn = os.path.splitext(fn)[0] + '-hdr.png'
|
|
shared.log.debug(f'Save: image="{fn}" type=PNG channels=16')
|
|
cv2.imwrite(fn, hdr)
|
|
# if opts.grid_save:
|
|
# images.save_image(grid, p.outpath_grids, "grid", p.seed, p.prompt, opts.grid_format, info=processed.info, grid=True, p=p)
|
|
if opts.return_grid:
|
|
grid = images.image_grid(imgs, rows=1)
|
|
imgs.append(grid)
|
|
|
|
processed = Processed(p, images_list=imgs, seed=p.seed, info=info)
|
|
return processed
|